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 { dtype: DataType, shape: Shape },
66    Sequence(Box<TypeProto>),
67    Optional(Box<TypeProto>),
68    Map { key: DataType, value: Box<TypeProto> },
69    SparseTensor { dtype: DataType, shape: Shape },
70}
71
72/// A reference to initializer (weight) data.
73///
74/// Small weights may be inlined; large weights are memory-mapped from an
75/// external file at load time (see `docs/ORT2.md` §12). The IR only stores the
76/// *reference*; the loader/`onnx-runtime-memory` crate performs the mmap.
77#[derive(Clone, Debug, PartialEq)]
78pub enum WeightRef {
79    /// Weight bytes stored inline in the model.
80    Inline(TensorData),
81    /// Weight bytes located in an external file at `[offset, offset+length)`.
82    External {
83        path: PathBuf,
84        offset: usize,
85        length: usize,
86        dtype: DataType,
87        dims: Vec<usize>,
88    },
89}
90
91impl WeightRef {
92    /// The element type of the referenced weight.
93    pub fn dtype(&self) -> DataType {
94        match self {
95            WeightRef::Inline(t) => t.dtype,
96            WeightRef::External { dtype, .. } => *dtype,
97        }
98    }
99
100    /// The static dimensions of the referenced weight.
101    pub fn dims(&self) -> &[usize] {
102        match self {
103            WeightRef::Inline(t) => &t.dims,
104            WeightRef::External { dims, .. } => dims,
105        }
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn tensor_numel_and_bytes() {
115        let t = TensorData::from_raw(DataType::Float32, vec![2, 3], vec![0u8; 24]);
116        assert_eq!(t.numel(), 6);
117        assert_eq!(t.expected_bytes(), 24);
118    }
119
120    #[test]
121    fn sub_byte_expected_bytes() {
122        let t = TensorData::from_raw(DataType::Int4, vec![3], vec![0u8; 2]);
123        assert_eq!(t.numel(), 3);
124        assert_eq!(t.expected_bytes(), 2); // 3 packed nibbles -> 2 bytes
125    }
126
127    #[test]
128    fn weight_ref_accessors() {
129        let w = WeightRef::External {
130            path: PathBuf::from("weights.bin"),
131            offset: 128,
132            length: 4096,
133            dtype: DataType::Float16,
134            dims: vec![64, 32],
135        };
136        assert_eq!(w.dtype(), DataType::Float16);
137        assert_eq!(w.dims(), &[64, 32]);
138    }
139}
onnx_runtime_ir/tensor.rs

onnx_runtime_ir/
tensor.rs
