burn_ir/
scalar.rs

1use burn_backend::{DType, Scalar};
2use burn_backend::{Element, ElementConversion};
3use core::hash::Hash;
4use serde::{Deserialize, Serialize};
5
6/// A scalar representation.
7#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
8#[allow(missing_docs)]
9pub enum ScalarIr {
10    Float(f64),
11    Int(i64),
12    UInt(u64),
13    Bool(bool),
14}
15
16impl Hash for ScalarIr {
17    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
18        match self {
19            ScalarIr::Float(x) => x.to_bits().hash(state),
20            ScalarIr::Int(x) => x.hash(state),
21            ScalarIr::UInt(x) => x.hash(state),
22            ScalarIr::Bool(x) => x.hash(state),
23        }
24    }
25}
26
27impl ScalarIr {
28    /// Creates a scalar with the specified data type.
29    pub fn new<E: ElementConversion>(value: E, dtype: &DType) -> Self {
30        if dtype.is_float() {
31            Self::Float(value.elem())
32        } else if dtype.is_int() {
33            Self::Int(value.elem())
34        } else if dtype.is_uint() {
35            Self::UInt(value.elem())
36        } else if dtype.is_bool() {
37            Self::Bool(value.elem())
38        } else {
39            unimplemented!("Scalar not supported for {dtype:?}")
40        }
41    }
42
43    /// Converts and returns the converted element.
44    pub fn elem<E: Element>(self) -> E {
45        match self {
46            ScalarIr::Float(x) => x.elem(),
47            ScalarIr::Int(x) => x.elem(),
48            ScalarIr::UInt(x) => x.elem(),
49            ScalarIr::Bool(x) => x.elem(),
50        }
51    }
52}
53
54// The enums are similar, but both types have different roles:
55// - `Scalar`: runtime literal value
56// - `ScalarIr`: serializable literal representation (used for IR)
57impl From<Scalar> for ScalarIr {
58    fn from(value: Scalar) -> Self {
59        match value {
60            Scalar::Float(x) => Self::Float(x),
61            Scalar::Int(x) => Self::Int(x),
62            Scalar::UInt(x) => Self::UInt(x),
63            Scalar::Bool(x) => Self::Bool(x),
64        }
65    }
66}
67
68impl From<ScalarIr> for Scalar {
69    fn from(value: ScalarIr) -> Self {
70        match value {
71            ScalarIr::Float(x) => Self::Float(x),
72            ScalarIr::Int(x) => Self::Int(x),
73            ScalarIr::UInt(x) => Self::UInt(x),
74            ScalarIr::Bool(x) => Self::Bool(x),
75        }
76    }
77}
burn_ir/scalar.rs

burn_ir/
scalar.rs
