onnx-ir 0.21.0

ONNX-IR is a pure Rust library for parsing ONNX models into an intermediate representation that can be used to generate code for various ML/DL frameworks
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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

//! Extension trait for burn_tensor::TensorData
//!
//! This module provides ONNX-specific helper methods for working with TensorData,
//! including type conversions and scalar extraction.

use burn_tensor::DType;

/// Re-export burn-tensor's TensorData for direct use
pub use burn_tensor::TensorData;

/// Extension trait providing ONNX-specific helper methods for TensorData
pub trait TensorDataExt {
    /// Alias for dtype for backward compatibility with ONNX naming
    fn elem_type(&self) -> DType;

    /// Extract the first element as f64, converting from any numeric type
    ///
    /// Useful for extracting scalar parameters from constant nodes.
    fn scalar_f64(&self) -> Result<f64, burn_tensor::DataError>;

    /// Extract the first element as f32, converting from any numeric type
    fn scalar_f32(&self) -> Result<f32, burn_tensor::DataError>;

    /// Extract the first element as i64, converting from any numeric type
    fn scalar_i64(&self) -> Result<i64, burn_tensor::DataError>;

    /// Convert to `Vec<i64>`, handling Int32 and Int64 types
    ///
    /// Useful for extracting indices, shapes, or other integer arrays that need to be i64.
    fn to_i64_vec(&self) -> Result<Vec<i64>, burn_tensor::DataError>;

    /// Convert to `Vec<f32>`, handling all numeric types with automatic conversion
    ///
    /// Useful for extracting numeric arrays that need to be f32.
    fn to_f32_vec(&self) -> Result<Vec<f32>, burn_tensor::DataError>;

    /// Convert to `Vec<f64>`, handling all numeric types with automatic conversion
    ///
    /// Useful when f64 precision is needed (e.g., exact comparisons against identity values).
    fn to_f64_vec(&self) -> Result<Vec<f64>, burn_tensor::DataError>;
}

impl TensorDataExt for burn_tensor::TensorData {
    fn elem_type(&self) -> DType {
        self.dtype
    }

    fn scalar_f64(&self) -> Result<f64, burn_tensor::DataError> {
        use burn_tensor::DType;
        match self.dtype {
            DType::F16 => {
                let val = self.as_slice::<half::f16>()?[0];
                Ok(f32::from(val) as f64)
            }
            DType::F32 => Ok(self.as_slice::<f32>()?[0] as f64),
            DType::F64 => Ok(self.as_slice::<f64>()?[0]),
            DType::I32 => Ok(self.as_slice::<i32>()?[0] as f64),
            DType::I64 => Ok(self.as_slice::<i64>()?[0] as f64),
            DType::I8 => Ok(self.as_slice::<i8>()?[0] as f64),
            DType::U8 => Ok(self.as_slice::<u8>()?[0] as f64),
            DType::U16 => Ok(self.as_slice::<u16>()?[0] as f64),
            other => Err(burn_tensor::DataError::TypeMismatch(format!(
                "Cannot convert {:?} to f64",
                other
            ))),
        }
    }

    fn scalar_f32(&self) -> Result<f32, burn_tensor::DataError> {
        self.scalar_f64().map(|v| v as f32)
    }

    fn scalar_i64(&self) -> Result<i64, burn_tensor::DataError> {
        use burn_tensor::DType;
        match self.dtype {
            DType::I64 => Ok(self.as_slice::<i64>()?[0]),
            DType::I32 => Ok(self.as_slice::<i32>()?[0] as i64),
            DType::I8 => Ok(self.as_slice::<i8>()?[0] as i64),
            DType::U8 => Ok(self.as_slice::<u8>()?[0] as i64),
            DType::U16 => Ok(self.as_slice::<u16>()?[0] as i64),
            DType::F32 => Ok(self.as_slice::<f32>()?[0] as i64),
            DType::F64 => Ok(self.as_slice::<f64>()?[0] as i64),
            other => Err(burn_tensor::DataError::TypeMismatch(format!(
                "Cannot convert {:?} to i64",
                other
            ))),
        }
    }

    fn to_i64_vec(&self) -> Result<Vec<i64>, burn_tensor::DataError> {
        use burn_tensor::DType;
        match self.dtype {
            DType::I64 | DType::I32 => self.clone().convert::<i64>().to_vec(),
            other => Err(burn_tensor::DataError::TypeMismatch(format!(
                "Cannot convert {:?} to Vec<i64>",
                other
            ))),
        }
    }

    fn to_f32_vec(&self) -> Result<Vec<f32>, burn_tensor::DataError> {
        use burn_tensor::DType;
        match self.dtype {
            DType::F32
            | DType::F64
            | DType::F16
            | DType::I64
            | DType::I32
            | DType::I8
            | DType::U8
            | DType::U16 => self.clone().convert::<f32>().to_vec(),
            other => Err(burn_tensor::DataError::TypeMismatch(format!(
                "Cannot convert {:?} to Vec<f32>",
                other
            ))),
        }
    }

    fn to_f64_vec(&self) -> Result<Vec<f64>, burn_tensor::DataError> {
        use burn_tensor::DType;
        match self.dtype {
            DType::F32
            | DType::F64
            | DType::F16
            | DType::I64
            | DType::I32
            | DType::I8
            | DType::U8
            | DType::U16 => self.clone().convert::<f64>().to_vec(),
            other => Err(burn_tensor::DataError::TypeMismatch(format!(
                "Cannot convert {:?} to Vec<f64>",
                other
            ))),
        }
    }
}