onnx_runtime_ir/
tensor.rs

1//! Constant tensor storage, weight references, and ONNX type descriptors.
2
3use std::path::PathBuf;
4
5use crate::dtype::DataType;
6use crate::shape::Shape;
7
8/// A concrete constant tensor held inline (e.g. an attribute value or a small
9/// initializer). Element bytes are stored little-endian and densely packed.
10///
11/// Large model weights are referenced lazily via [`WeightRef`] instead.
12#[derive(Clone, Debug, PartialEq)]
13pub struct TensorData {
14    pub name: Option<String>,
15    pub dtype: DataType,
16    /// Static dimensions (constants always have a fully known shape).
17    pub dims: Vec<usize>,
18    /// Raw little-endian element bytes. Sub-byte values are densely packed
19    /// (two 4-bit or four 2-bit elements per byte); for [`DataType::String`]
20    /// this is empty and `strings` is used instead.
21    pub data: Vec<u8>,
22    /// String payloads for [`DataType::String`] tensors.
23    pub strings: Vec<String>,
24}
25
26impl TensorData {
27    /// A numeric tensor from raw little-endian bytes.
28    pub fn from_raw(dtype: DataType, dims: Vec<usize>, data: Vec<u8>) -> Self {
29        Self {
30            name: None,
31            dtype,
32            dims,
33            data,
34            strings: Vec::new(),
35        }
36    }
37
38    /// Number of elements (product of dims; `1` for a scalar).
39    pub fn numel(&self) -> usize {
40        self.dims.iter().product()
41    }
42
43    /// Expected byte length for `numel` elements of `dtype`, accounting for
44    /// sub-byte packing.
45    pub fn expected_bytes(&self) -> usize {
46        self.dtype.storage_bytes(self.numel())
47    }
48}
49
50/// A sparse constant tensor in COO form.
51#[derive(Clone, Debug, PartialEq)]
52pub struct SparseTensorData {
53    /// Non-zero values.
54    pub values: TensorData,
55    /// Indices of the non-zero values (int64), shape `[nnz, rank]` or `[nnz]`.
56    pub indices: TensorData,
57    /// Dense shape.
58    pub dims: Vec<usize>,
59}
60
61/// An ONNX `TypeProto`: the type of a value, which may be a tensor or a
62/// container of tensors (see `docs/ORT2.md` §3.2).
63#[derive(Clone, Debug, PartialEq)]
64pub enum TypeProto {
65    Tensor {
66        dtype: DataType,
67        shape: Shape,
68    },
69    Sequence(Box<TypeProto>),
70    Optional(Box<TypeProto>),
71    Map {
72        key: DataType,
73        value: Box<TypeProto>,
74    },
75    SparseTensor {
76        dtype: DataType,
77        shape: Shape,
78    },
79}
80
81/// A reference to initializer (weight) data.
82///
83/// Small weights may be inlined; large weights are memory-mapped from an
84/// external file at load time (see `docs/ORT2.md` §12). The IR only stores the
85/// *reference*; the loader/`onnx-runtime-memory` crate performs the mmap.
86#[derive(Clone, Debug, PartialEq)]
87pub enum WeightRef {
88    /// Weight bytes stored inline in the model.
89    Inline(TensorData),
90    /// Weight bytes located in an external file at `[offset, offset+length)`.
91    External {
92        path: PathBuf,
93        offset: usize,
94        length: usize,
95        dtype: DataType,
96        dims: Vec<usize>,
97    },
98}
99
100impl WeightRef {
101    /// The element type of the referenced weight.
102    pub fn dtype(&self) -> DataType {
103        match self {
104            WeightRef::Inline(t) => t.dtype,
105            WeightRef::External { dtype, .. } => *dtype,
106        }
107    }
108
109    /// The static dimensions of the referenced weight.
110    pub fn dims(&self) -> &[usize] {
111        match self {
112            WeightRef::Inline(t) => &t.dims,
113            WeightRef::External { dims, .. } => dims,
114        }
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    #[test]
123    fn tensor_numel_and_bytes() {
124        let t = TensorData::from_raw(DataType::Float32, vec![2, 3], vec![0u8; 24]);
125        assert_eq!(t.numel(), 6);
126        assert_eq!(t.expected_bytes(), 24);
127    }
128
129    #[test]
130    fn sub_byte_expected_bytes() {
131        let t = TensorData::from_raw(DataType::Int4, vec![3], vec![0u8; 2]);
132        assert_eq!(t.numel(), 3);
133        assert_eq!(t.expected_bytes(), 2); // 3 packed nibbles -> 2 bytes
134    }
135
136    #[test]
137    fn weight_ref_accessors() {
138        let w = WeightRef::External {
139            path: PathBuf::from("weights.bin"),
140            offset: 128,
141            length: 4096,
142            dtype: DataType::Float16,
143            dims: vec![64, 32],
144        };
145        assert_eq!(w.dtype(), DataType::Float16);
146        assert_eq!(w.dims(), &[64, 32]);
147    }
148}
onnx_runtime_ir/tensor.rs

onnx_runtime_ir/
tensor.rs
