nir_rs/
types.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! Shared value types for NIR tensors and metadata.
4//!
5//! Python NIR stores parameters as `numpy.ndarray`. This module provides a
6//! compact, owned representation suitable for an in-memory Rust IR. HDF5
7//! decoding (v0.3) will map wire arrays into [`Tensor`].
8
9use crate::error::{NirError, Result};
10
11/// Element type of a contiguous tensor buffer.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
13#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
14#[non_exhaustive]
15pub enum DType {
16    /// 32-bit IEEE floating point.
17    F32,
18    /// 64-bit IEEE floating point.
19    F64,
20    /// 64-bit signed integer.
21    I64,
22    /// Boolean.
23    Bool,
24}
25
26impl DType {
27    /// Size of one element in bytes.
28    #[must_use]
29    pub const fn size_of(self) -> usize {
30        match self {
31            Self::F32 => 4,
32            Self::F64 => 8,
33            Self::I64 => 8,
34            Self::Bool => 1,
35        }
36    }
37}
38
39/// Contiguous numeric payload backing a [`Tensor`].
40#[derive(Debug, Clone, PartialEq)]
41#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
42pub enum TensorData {
43    /// `f32` elements.
44    F32(Vec<f32>),
45    /// `f64` elements.
46    F64(Vec<f64>),
47    /// `i64` elements.
48    I64(Vec<i64>),
49    /// `bool` elements.
50    Bool(Vec<bool>),
51}
52
53impl TensorData {
54    /// Number of elements.
55    #[must_use]
56    pub fn len(&self) -> usize {
57        match self {
58            Self::F32(v) => v.len(),
59            Self::F64(v) => v.len(),
60            Self::I64(v) => v.len(),
61            Self::Bool(v) => v.len(),
62        }
63    }
64
65    /// Whether the buffer is empty.
66    #[must_use]
67    pub fn is_empty(&self) -> bool {
68        self.len() == 0
69    }
70
71    /// Element dtype of this payload.
72    #[must_use]
73    pub const fn dtype(&self) -> DType {
74        match self {
75            Self::F32(_) => DType::F32,
76            Self::F64(_) => DType::F64,
77            Self::I64(_) => DType::I64,
78            Self::Bool(_) => DType::Bool,
79        }
80    }
81}
82
83/// Dense, row-major tensor (shape + contiguous typed data).
84///
85/// Shape is listed outer-to-inner (C-order), matching typical NumPy layout.
86///
87/// Element dtype is always derived from the payload via [`Tensor::dtype`].
88///
89/// Shape and data are **private** so callers cannot break the
90/// `shape product == data.len()` invariant after construction. Use
91/// [`Tensor::shape`] / [`Tensor::data`] accessors.
92///
93/// # PartialEq
94///
95/// Equality is exact element-wise (IEEE). In particular, `NaN != NaN`, matching
96/// Rust's default float `PartialEq`.
97///
98/// # Serde
99///
100/// With the `serde` feature, tensors encode as `{ "shape": ..., "data": ... }`.
101/// Deserialization always calls [`Tensor::new`], preserving the private-field
102/// invariant. JSON is debug-only and cannot faithfully represent non-finite
103/// floats; HDF5 `.nir` remains the NIR interchange format.
104#[derive(Debug, Clone, PartialEq)]
105pub struct Tensor {
106    /// Axis lengths.
107    shape: Vec<usize>,
108    /// Contiguous elements in C-order.
109    data: TensorData,
110}
111
112#[cfg(feature = "serde")]
113impl serde::Serialize for Tensor {
114    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
115    where
116        S: serde::Serializer,
117    {
118        #[derive(serde::Serialize)]
119        struct TensorRef<'a> {
120            shape: &'a [usize],
121            data: &'a TensorData,
122        }
123
124        serde::Serialize::serialize(
125            &TensorRef {
126                shape: self.shape(),
127                data: self.data(),
128            },
129            serializer,
130        )
131    }
132}
133
134#[cfg(feature = "serde")]
135impl<'de> serde::Deserialize<'de> for Tensor {
136    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
137    where
138        D: serde::Deserializer<'de>,
139    {
140        #[derive(serde::Deserialize)]
141        struct TensorOwned {
142            shape: Vec<usize>,
143            data: TensorData,
144        }
145
146        let tensor = <TensorOwned as serde::Deserialize>::deserialize(deserializer)?;
147        Self::new(tensor.shape, tensor.data).map_err(serde::de::Error::custom)
148    }
149}
150
151impl Tensor {
152    /// Build a tensor from shape and typed data, checking length against shape.
153    pub fn new(shape: impl Into<Vec<usize>>, data: TensorData) -> Result<Self> {
154        let shape = shape.into();
155        check_shape_len(&shape, data.len())?;
156        Ok(Self { shape, data })
157    }
158
159    /// Element dtype of this tensor (derived from the payload).
160    #[must_use]
161    pub const fn dtype(&self) -> DType {
162        self.data.dtype()
163    }
164
165    /// Axis lengths (outer-to-inner / C-order).
166    #[must_use]
167    pub fn shape(&self) -> &[usize] {
168        &self.shape
169    }
170
171    /// Contiguous payload.
172    #[must_use]
173    pub fn data(&self) -> &TensorData {
174        &self.data
175    }
176
177    /// Take the payload, consuming the tensor.
178    ///
179    /// Crate-internal: the HDF5 reader decodes integer wire fields through a
180    /// `Tensor` and then wants the `Vec` itself. Cloning out of [`data`] would
181    /// hold both buffers live at once, doubling peak memory on every `shape`,
182    /// `stride` and `dilation` read for no benefit.
183    ///
184    /// [`data`]: Self::data
185    #[must_use]
186    #[cfg(feature = "hdf5")]
187    pub(crate) fn into_data(self) -> TensorData {
188        self.data
189    }
190
191    /// `f32` tensor; `data.len()` must equal the product of `shape`.
192    pub fn from_f32(shape: impl Into<Vec<usize>>, data: impl Into<Vec<f32>>) -> Result<Self> {
193        Self::new(shape, TensorData::F32(data.into()))
194    }
195
196    /// `f64` tensor; `data.len()` must equal the product of `shape`.
197    pub fn from_f64(shape: impl Into<Vec<usize>>, data: impl Into<Vec<f64>>) -> Result<Self> {
198        Self::new(shape, TensorData::F64(data.into()))
199    }
200
201    /// `i64` tensor; `data.len()` must equal the product of `shape`.
202    pub fn from_i64(shape: impl Into<Vec<usize>>, data: impl Into<Vec<i64>>) -> Result<Self> {
203        Self::new(shape, TensorData::I64(data.into()))
204    }
205
206    /// `bool` tensor; `data.len()` must equal the product of `shape`.
207    pub fn from_bool(shape: impl Into<Vec<usize>>, data: impl Into<Vec<bool>>) -> Result<Self> {
208        Self::new(shape, TensorData::Bool(data.into()))
209    }
210
211    /// Rank-0 (scalar) `f32` tensor.
212    #[must_use]
213    pub fn scalar_f32(value: f32) -> Self {
214        Self {
215            shape: vec![],
216            data: TensorData::F32(vec![value]),
217        }
218    }
219
220    /// Rank-0 (scalar) `f64` tensor.
221    #[must_use]
222    pub fn scalar_f64(value: f64) -> Self {
223        Self {
224            shape: vec![],
225            data: TensorData::F64(vec![value]),
226        }
227    }
228
229    /// Rank-0 (scalar) `i64` tensor.
230    #[must_use]
231    pub fn scalar_i64(value: i64) -> Self {
232        Self {
233            shape: vec![],
234            data: TensorData::I64(vec![value]),
235        }
236    }
237
238    /// A tensor of zeros with the same shape and dtype as `self`.
239    ///
240    /// Mirrors `numpy.zeros_like`, which upstream NIR uses to default absent
241    /// optional wire fields (`v_reset`). [`DType::Bool`] zeroes to `false`.
242    #[must_use]
243    pub fn zeros_like(&self) -> Self {
244        let n = self.data.len();
245        let data = match self.dtype() {
246            DType::F32 => TensorData::F32(vec![0.0; n]),
247            DType::F64 => TensorData::F64(vec![0.0; n]),
248            DType::I64 => TensorData::I64(vec![0; n]),
249            DType::Bool => TensorData::Bool(vec![false; n]),
250        };
251        Self {
252            shape: self.shape.clone(),
253            data,
254        }
255    }
256
257    /// A tensor of ones with the same shape and dtype as `self`.
258    ///
259    /// Mirrors `numpy.ones_like`, which upstream NIR uses to default an absent
260    /// `w_in`. [`DType::Bool`] ones to `true`.
261    #[must_use]
262    pub fn ones_like(&self) -> Self {
263        let n = self.data.len();
264        let data = match self.dtype() {
265            DType::F32 => TensorData::F32(vec![1.0; n]),
266            DType::F64 => TensorData::F64(vec![1.0; n]),
267            DType::I64 => TensorData::I64(vec![1; n]),
268            DType::Bool => TensorData::Bool(vec![true; n]),
269        };
270        Self {
271            shape: self.shape.clone(),
272            data,
273        }
274    }
275
276    /// Number of elements implied by `shape` (empty shape → 1).
277    #[must_use]
278    pub fn numel(&self) -> usize {
279        // Invariant: construction ensures product fits and matches data.len().
280        shape_product(&self.shape).expect("tensor shape product overflow")
281    }
282
283    /// Number of dimensions.
284    #[must_use]
285    pub fn ndim(&self) -> usize {
286        self.shape.len()
287    }
288}
289
290/// Free-form metadata map (Python NIR `metadata: Dict[str, Any]`).
291pub type MetadataMap = std::collections::HashMap<String, MetadataValue>;
292
293/// Free-form metadata values attached to graphs and nodes.
294#[derive(Debug, Clone, PartialEq)]
295#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
296#[non_exhaustive]
297pub enum MetadataValue {
298    /// UTF-8 string.
299    String(String),
300    /// List of UTF-8 strings.
301    ///
302    /// Python's `metadata: Dict[str, Any]` admits a `list[str]`, which h5py
303    /// stores as a multi-element string dataset. Without this variant such a
304    /// file cannot be decoded at all, since [`Tensor`] carries only numeric
305    /// and boolean payloads.
306    StringList(Vec<String>),
307    /// 64-bit float.
308    F64(f64),
309    /// 64-bit signed integer.
310    I64(i64),
311    /// Boolean.
312    Bool(bool),
313    /// Dense tensor.
314    Tensor(Tensor),
315}
316
317/// Product of shape dims; empty shape is a scalar (1 element).
318/// Returns `None` if the product overflows `usize`.
319fn shape_product(shape: &[usize]) -> Option<usize> {
320    if shape.is_empty() {
321        Some(1)
322    } else {
323        shape.iter().try_fold(1usize, |acc, &d| acc.checked_mul(d))
324    }
325}
326
327fn check_shape_len(shape: &[usize], len: usize) -> Result<()> {
328    let expected = shape_product(shape).ok_or_else(|| {
329        NirError::InvalidTensor(format!("shape product overflows usize (shape={shape:?})"))
330    })?;
331    if expected != len {
332        return Err(NirError::InvalidTensor(format!(
333            "shape product {expected} != data len {len} (shape={shape:?})"
334        )));
335    }
336    Ok(())
337}
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342
343    #[test]
344    fn from_f32_ok() {
345        let t = Tensor::from_f32(vec![2, 3], vec![1., 2., 3., 4., 5., 6.]).unwrap();
346        assert_eq!(t.dtype(), DType::F32);
347        assert_eq!(t.numel(), 6);
348        assert_eq!(t.ndim(), 2);
349    }
350
351    #[test]
352    fn from_f64_ok() {
353        let t = Tensor::from_f64([2], vec![1.0, 2.0]).unwrap();
354        assert_eq!(t.dtype(), DType::F64);
355        assert_eq!(t.numel(), 2);
356    }
357
358    #[test]
359    fn length_mismatch_f32() {
360        let err = Tensor::from_f32(vec![2, 2], vec![1., 2., 3.]).unwrap_err();
361        assert!(matches!(err, NirError::InvalidTensor(_)));
362        assert!(err.to_string().contains("shape product 4 != data len 3"));
363    }
364
365    #[test]
366    fn length_mismatch_f64() {
367        let err = Tensor::from_f64(vec![3], vec![1.0]).unwrap_err();
368        assert!(matches!(err, NirError::InvalidTensor(_)));
369    }
370
371    #[test]
372    fn scalar_has_empty_shape_one_element() {
373        let t = Tensor::scalar_f64(0.5);
374        assert!(t.shape().is_empty());
375        assert_eq!(t.numel(), 1);
376        assert_eq!(t.data().len(), 1);
377    }
378
379    #[test]
380    fn empty_shape_rejects_wrong_len() {
381        let err = Tensor::from_f32(Vec::<usize>::new(), vec![1., 2.]).unwrap_err();
382        assert!(matches!(err, NirError::InvalidTensor(_)));
383    }
384
385    #[test]
386    fn i64_and_bool_constructors() {
387        let i = Tensor::from_i64([2], vec![1, 2]).unwrap();
388        assert_eq!(i.dtype(), DType::I64);
389        let b = Tensor::from_bool([2], vec![true, false]).unwrap();
390        assert_eq!(b.dtype(), DType::Bool);
391    }
392
393    #[test]
394    fn dtype_size_of() {
395        assert_eq!(DType::F32.size_of(), 4);
396        assert_eq!(DType::F64.size_of(), 8);
397        assert_eq!(DType::I64.size_of(), 8);
398        assert_eq!(DType::Bool.size_of(), 1);
399    }
400
401    #[test]
402    fn shape_product_overflow_rejected() {
403        let err = Tensor::from_f32(vec![usize::MAX, usize::MAX], vec![1.0]).unwrap_err();
404        assert!(matches!(err, NirError::InvalidTensor(_)));
405        assert!(err.to_string().contains("overflows"));
406    }
407
408    #[test]
409    fn zeros_like_preserves_shape_and_dtype() {
410        let t = Tensor::from_f32(vec![2, 2], vec![1., 2., 3., 4.]).unwrap();
411        let z = t.zeros_like();
412        assert_eq!(z.shape(), t.shape());
413        assert_eq!(z.dtype(), DType::F32);
414        assert_eq!(z.data(), &TensorData::F32(vec![0.0; 4]));
415    }
416
417    #[test]
418    fn ones_like_preserves_shape_and_dtype() {
419        let t = Tensor::from_f64(vec![3], vec![7.0, 8.0, 9.0]).unwrap();
420        let o = t.ones_like();
421        assert_eq!(o.shape(), [3]);
422        assert_eq!(o.data(), &TensorData::F64(vec![1.0; 3]));
423    }
424
425    #[test]
426    fn zeros_and_ones_like_cover_int_and_bool() {
427        let i = Tensor::from_i64([2], vec![5, 6]).unwrap();
428        assert_eq!(i.zeros_like().data(), &TensorData::I64(vec![0, 0]));
429        assert_eq!(i.ones_like().data(), &TensorData::I64(vec![1, 1]));
430
431        let b = Tensor::from_bool([2], vec![true, false]).unwrap();
432        assert_eq!(b.zeros_like().data(), &TensorData::Bool(vec![false, false]));
433        assert_eq!(b.ones_like().data(), &TensorData::Bool(vec![true, true]));
434    }
435
436    #[test]
437    fn zeros_like_of_scalar_is_scalar() {
438        let z = Tensor::scalar_f64(3.5).zeros_like();
439        assert!(z.shape().is_empty());
440        assert_eq!(z.numel(), 1);
441    }
442
443    #[test]
444    fn metadata_variants() {
445        let m = MetadataValue::String("note".into());
446        assert!(matches!(m, MetadataValue::String(_)));
447        let _ = MetadataValue::F64(1.0);
448        let _ = MetadataValue::I64(2);
449        let _ = MetadataValue::Bool(true);
450        let _ = MetadataValue::Tensor(Tensor::scalar_f32(0.0));
451    }
452}
nir_rs/types.rs

nir_rs/
types.rs
