burn-ir 0.21.0

Intermediate representation for the Burn framework
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

use burn_backend::{DType, Scalar};
use burn_backend::{Element, ElementConversion};
use core::hash::Hash;
use serde::{Deserialize, Serialize};

/// A scalar representation.
#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
#[allow(missing_docs)]
pub enum ScalarIr {
    Float(f64),
    Int(i64),
    UInt(u64),
    Bool(bool),
}

impl Hash for ScalarIr {
    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
        match self {
            ScalarIr::Float(x) => x.to_bits().hash(state),
            ScalarIr::Int(x) => x.hash(state),
            ScalarIr::UInt(x) => x.hash(state),
            ScalarIr::Bool(x) => x.hash(state),
        }
    }
}

impl ScalarIr {
    /// Creates a scalar with the specified data type.
    pub fn new<E: ElementConversion>(value: E, dtype: &DType) -> Self {
        if dtype.is_float() {
            Self::Float(value.elem())
        } else if dtype.is_int() {
            Self::Int(value.elem())
        } else if dtype.is_uint() {
            Self::UInt(value.elem())
        } else if dtype.is_bool() {
            Self::Bool(value.elem())
        } else {
            unimplemented!("Scalar not supported for {dtype:?}")
        }
    }

    /// Converts and returns the converted element.
    pub fn elem<E: Element>(self) -> E {
        match self {
            ScalarIr::Float(x) => x.elem(),
            ScalarIr::Int(x) => x.elem(),
            ScalarIr::UInt(x) => x.elem(),
            ScalarIr::Bool(x) => x.elem(),
        }
    }
}

// The enums are similar, but both types have different roles:
// - `Scalar`: runtime literal value
// - `ScalarIr`: serializable literal representation (used for IR)
impl From<Scalar> for ScalarIr {
    fn from(value: Scalar) -> Self {
        match value {
            Scalar::Float(x) => Self::Float(x),
            Scalar::Int(x) => Self::Int(x),
            Scalar::UInt(x) => Self::UInt(x),
            Scalar::Bool(x) => Self::Bool(x),
        }
    }
}

impl From<ScalarIr> for Scalar {
    fn from(value: ScalarIr) -> Self {
        match value {
            ScalarIr::Float(x) => Self::Float(x),
            ScalarIr::Int(x) => Self::Int(x),
            ScalarIr::UInt(x) => Self::UInt(x),
            ScalarIr::Bool(x) => Self::Bool(x),
        }
    }
}