burn_std/data/tensor/
base.rs

1use core::f32;
2
3use alloc::format;
4use alloc::vec::Vec;
5use bytemuck::checked::CheckedCastError;
6use rand::Rng;
7use thiserror::Error;
8
9use crate::Scalar;
10use crate::distribution::Distribution;
11use crate::element::{Element, ElementConversion};
12use crate::tensor::DType;
13use crate::{
14    AccessError, BoolStore, Bytes, ExecutionError, QuantMode, QuantScheme, QuantValue,
15    QuantizedBytes, Reader, Shape, Writer, bf16, f16, flex32,
16};
17
18use serde::{Deserialize, Serialize};
19
20/// Errors that can occur while accessing or converting [`TensorData`].
21#[derive(Debug, Error, PartialEq, Eq)]
22pub enum DataError {
23    /// Host access to the underlying storage failed.
24    #[error("Failed to access TensorData storage: {0}")]
25    StorageAccess(#[from] AccessError),
26
27    /// The stored bytes aren't a valid representation of the requested element type.
28    #[error("TensorData storage is invalid for the requested element type: {0}")]
29    InvalidRepresentation(CheckedCastError),
30
31    /// The stored dtype doesn't match the requested dtype.
32    #[error("Expected data type {expected:?}, but got {actual:?}")]
33    DTypeMismatch {
34        /// The expected storage DType.
35        expected: DType,
36
37        /// The actual storage DType.
38        actual: DType,
39    },
40
41    /// Unsupported data conversion.
42    #[error("Unsupported data conversion from {from:?} to {to:?}")]
43    UnsupportedConversion {
44        /// The source DType.
45        from: DType,
46
47        /// The destination DType.
48        to: DType,
49    },
50
51    /// The byte storage doesn't match the number of elements described by the shape.
52    #[error("TensorData shape describes {expected} element(s), but storage contains {actual}")]
53    ElementCountMismatch {
54        /// The number of elements described by the shape.
55        expected: usize,
56
57        /// The number of elements present in storage.
58        actual: usize,
59    },
60
61    /// The byte storage length doesn't match the size implied by the shape and dtype.
62    #[error(
63        "Shape {:?} with dtype {dtype:?} is invalid for input of size {len} bytes",
64        shape.as_slice()
65    )]
66    InvalidByteLength {
67        /// The shape describing the tensor.
68        shape: Shape,
69
70        /// The storage DType.
71        dtype: DType,
72
73        /// The number of bytes present in storage.
74        len: usize,
75    },
76
77    /// The quantized storage layout is invalid or cannot be represented for the tensor shape.
78    #[error("Quantized layout for shape {:?} and dtype {dtype:?} is invalid", shape.as_slice())]
79    InvalidQuantizedLayout {
80        /// The shape describing the tensor.
81        shape: Shape,
82
83        /// The quantized storage dtype.
84        dtype: DType,
85    },
86}
87
88/// Errors that can occur while reading host data from a tensor.
89#[derive(Debug, Error)]
90pub enum TensorReadError {
91    /// Tensor execution failed while reading the data from the device.
92    #[error(transparent)]
93    Execution(#[from] ExecutionError),
94
95    /// The resulting [`TensorData`] could not satisfy the requested data operation.
96    #[error(transparent)]
97    Data(#[from] DataError),
98
99    /// The tensor shape does not satisfy the read operation's element-count requirement.
100    #[error("Expected {expected} tensor element(s), but got {actual}")]
101    InvalidShape {
102        /// The required number of elements.
103        expected: usize,
104
105        /// The actual number of elements.
106        actual: usize,
107    },
108}
109
110impl DataError {
111    /// Creates a [`DataError::DTypeMismatch`] for the requested element type `E`.
112    pub(super) fn dtype_mismatch_as<E: Element>(actual: DType) -> Self {
113        Self::DTypeMismatch {
114            expected: E::dtype(),
115            actual,
116        }
117    }
118}
119
120/// Data structure for tensors.
121///
122/// The byte length of the storage always matches the size implied by the shape and dtype. Element
123/// values aren't validated on construction; the checked accessors reject invalid representations.
124#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
125#[serde(try_from = "TensorDataDe")]
126pub struct TensorData {
127    /// The values of the tensor (as bytes).
128    pub(in crate::data) bytes: Bytes,
129
130    /// The shape of the tensor.
131    #[serde(with = "shape_inner")]
132    pub(in crate::data) shape: Shape,
133
134    /// The data type of the tensor.
135    pub(in crate::data) dtype: DType,
136}
137
138/// Unvalidated wire form of [`TensorData`], used so that deserialization enforces the same
139/// byte-length invariant as the constructors instead of trusting the encoded shape.
140#[derive(Deserialize)]
141struct TensorDataDe {
142    bytes: Bytes,
143    #[serde(with = "shape_inner")]
144    shape: Shape,
145    dtype: DType,
146}
147
148impl TryFrom<TensorDataDe> for TensorData {
149    type Error = DataError;
150
151    fn try_from(data: TensorDataDe) -> Result<Self, Self::Error> {
152        TensorData::try_from_bytes(data.bytes, data.shape, data.dtype)
153    }
154}
155
156// For backward compatibility with shape `Vec<usize>`
157mod shape_inner {
158    use crate::SmallVec;
159
160    use super::*;
161
162    pub fn serialize<S: serde::Serializer>(
163        shape: &Shape,
164        serializer: S,
165    ) -> Result<S::Ok, S::Error> {
166        shape.as_slice().serialize(serializer)
167    }
168
169    pub fn deserialize<'de, D: serde::Deserializer<'de>>(
170        deserializer: D,
171    ) -> Result<Shape, D::Error> {
172        let dims = SmallVec::<[usize; _]>::deserialize(deserializer)?;
173        Ok(Shape::new_raw(dims))
174    }
175}
176
177impl TensorData {
178    /// Creates a new tensor data structure.
179    pub fn new<E: Element, S: Into<Shape>>(value: Vec<E>, shape: S) -> Self {
180        // Ensure shape is valid
181        let shape = shape.into();
182        Self::check_data_len(&value, &shape);
183
184        Self {
185            bytes: Bytes::from_elems(value),
186            shape,
187            dtype: E::dtype(),
188        }
189    }
190
191    /// Creates a new quantized tensor data structure.
192    pub fn quantized<E: Element, S: Into<Shape>>(
193        value: Vec<E>,
194        shape: S,
195        scheme: QuantScheme,
196        qparams: &[f32],
197        global: Option<f32>,
198    ) -> Self {
199        let shape = shape.into();
200        Self::check_data_len(&value, &shape);
201
202        let q_bytes = QuantizedBytes::new(value, shape.clone(), scheme, qparams, global);
203
204        Self {
205            bytes: q_bytes.bytes,
206            shape,
207            dtype: DType::QFloat(q_bytes.scheme),
208        }
209    }
210
211    /// Creates a new tensor data structure from raw bytes.
212    ///
213    /// # Panics
214    ///
215    /// Panics if the byte length doesn't match the shape and dtype, or the quantized scheme is
216    /// invalid for the shape. See [`TensorData::try_from_bytes`] for the fallible version.
217    #[track_caller]
218    pub fn from_bytes<S: Into<Shape>>(bytes: Bytes, shape: S, dtype: DType) -> Self {
219        match Self::try_from_bytes(bytes, shape, dtype) {
220            Ok(data) => data,
221            Err(err) => panic!("{err}"),
222        }
223    }
224
225    /// Creates a new tensor data structure from raw bytes stored in a vector.
226    ///
227    /// Prefer [`TensorData::new`] or [`TensorData::quantized`] over this method unless you are
228    /// certain that the element values are valid for `dtype`; only the length is checked.
229    ///
230    /// # Panics
231    ///
232    /// Panics if the byte length doesn't match the shape and dtype, or the quantized scheme is
233    /// invalid for the shape. See [`TensorData::try_from_bytes_vec`] for the fallible version.
234    #[track_caller]
235    pub fn from_bytes_vec<S: Into<Shape>>(bytes: Vec<u8>, shape: S, dtype: DType) -> Self {
236        Self::from_bytes(Bytes::from_bytes_vec(bytes), shape, dtype)
237    }
238
239    /// Creates a new tensor data structure from raw bytes.
240    ///
241    /// # Errors
242    ///
243    /// Returns [`DataError::InvalidByteLength`] if the byte length doesn't match the storage
244    /// layout described by the shape and dtype, including quantized schemes.
245    /// Returns [`DataError::InvalidQuantizedLayout`] if a quantized storage layout is invalid for
246    /// the shape or its size overflows.
247    ///
248    /// The element values themselves aren't validated here (e.g. a `bool` byte other than 0 or
249    /// 1); the checked accessors ([`TensorData::as_slice`], [`TensorData::try_to_vec`],
250    /// [`TensorData::try_view`]) reject invalid representations.
251    pub fn try_from_bytes<S: Into<Shape>>(
252        bytes: Bytes,
253        shape: S,
254        dtype: DType,
255    ) -> Result<Self, DataError> {
256        let shape = shape.into();
257        // `dtype.size()` is not the stored width of quantized data (packed values plus appended
258        // scales), so quantized dtypes use their storage scheme to determine the byte length.
259        // Products are checked because untrusted shapes and schemes can overflow `usize`.
260        let expected = match dtype {
261            DType::QFloat(scheme) => {
262                match crate::tensor::quantization::try_quantized_data_len(&scheme, &shape) {
263                    Some(expected) => Some(expected),
264                    None => return Err(DataError::InvalidQuantizedLayout { shape, dtype }),
265                }
266            }
267            _ => checked_numel(&shape).and_then(|numel| numel.checked_mul(dtype.size())),
268        };
269
270        if expected != Some(bytes.len()) {
271            return Err(DataError::InvalidByteLength {
272                shape,
273                dtype,
274                len: bytes.len(),
275            });
276        }
277
278        Ok(Self {
279            bytes,
280            shape,
281            dtype,
282        })
283    }
284
285    /// Creates a new tensor data structure from raw bytes stored in a vector.
286    ///
287    /// # Errors
288    ///
289    /// See [`TensorData::try_from_bytes`].
290    pub fn try_from_bytes_vec<S: Into<Shape>>(
291        bytes: Vec<u8>,
292        shape: S,
293        dtype: DType,
294    ) -> Result<Self, DataError> {
295        Self::try_from_bytes(Bytes::from_bytes_vec(bytes), shape, dtype)
296    }
297
298    /// Builds tensor data without validating it, so tests can exercise the defensive checks
299    /// on storage that the public constructors reject.
300    #[cfg(test)]
301    pub(crate) fn from_bytes_unchecked<S: Into<Shape>>(
302        bytes: Bytes,
303        shape: S,
304        dtype: DType,
305    ) -> Self {
306        Self {
307            bytes,
308            shape: shape.into(),
309            dtype,
310        }
311    }
312
313    /// Returns the shape of the tensor data.
314    pub fn shape(&self) -> &Shape {
315        &self.shape
316    }
317
318    /// Returns the data type of the tensor data.
319    pub fn dtype(&self) -> DType {
320        self.dtype
321    }
322
323    /// Returns the underlying byte storage.
324    pub fn bytes(&self) -> &Bytes {
325        &self.bytes
326    }
327
328    /// Consumes the tensor data and returns its byte storage, shape and data type.
329    pub fn into_parts(self) -> (Bytes, Shape, DType) {
330        (self.bytes, self.shape, self.dtype)
331    }
332
333    /// Gives `f` mutable access to the byte storage of every item at once, e.g. to move it into
334    /// staging memory in a single batch.
335    ///
336    /// # Panics
337    ///
338    /// Panics if `f` changes the length of any storage. Items whose length changed are replaced
339    /// with empty data first, and the same happens if `f` panics, so no item is left with a
340    /// length that disagrees with its shape.
341    #[track_caller]
342    pub fn with_bytes_mut<'a, R>(
343        data: impl IntoIterator<Item = &'a mut TensorData>,
344        f: impl FnOnce(Vec<&mut Bytes>) -> R,
345    ) -> R {
346        let mut guard = LengthGuard::new(data.into_iter().collect());
347        let out = f(guard.data.iter_mut().map(|data| &mut data.bytes).collect());
348        let changed = guard.reset_changed();
349        assert!(!changed, "TensorData byte storage length must not change");
350        out
351    }
352
353    // Check that the input vector contains a correct number of elements
354    fn check_data_len<E: Element>(data: &[E], shape: &Shape) {
355        let expected_data_len = numel(shape);
356        let num_data = data.len();
357        assert_eq!(
358            expected_data_len, num_data,
359            "Shape {shape:?} is invalid for input of size {num_data:?}",
360        );
361    }
362
363    /// Returns the immutable slice view of the tensor data.
364    ///
365    /// This materializes lazy storage into host-accessible memory when necessary.
366    ///
367    /// # Errors
368    ///
369    /// Returns an error if host access fails, the target element type doesn't match the stored
370    /// type, or the stored byte representation is invalid for `E`.
371    pub fn as_slice<E: Element>(&self) -> Result<&[E], DataError> {
372        if self.matches_target_dtype::<E>() {
373            bytemuck::checked::try_cast_slice(self.bytes.read(Reader::new())?)
374                .map_err(DataError::InvalidRepresentation)
375        } else {
376            Err(DataError::dtype_mismatch_as::<E>(self.dtype))
377        }
378    }
379
380    /// Returns the mutable slice view of the tensor data.
381    ///
382    /// This materializes lazy storage and performs copy-on-write when necessary.
383    ///
384    /// # Errors
385    ///
386    /// Returns an error if host access fails, the target element type doesn't match the stored
387    /// type, or the stored byte representation is invalid for `E`.
388    pub fn as_mut_slice<E: Element>(&mut self) -> Result<&mut [E], DataError> {
389        if self.matches_target_dtype::<E>() {
390            bytemuck::checked::try_cast_slice_mut(self.bytes.write(Writer::new())?)
391                .map_err(DataError::InvalidRepresentation)
392        } else {
393            Err(DataError::dtype_mismatch_as::<E>(self.dtype))
394        }
395    }
396
397    pub(super) fn matches_target_dtype<E: Element>(&self) -> bool {
398        let target_dtype = E::dtype();
399        match self.dtype {
400            DType::Bool(BoolStore::U8) => {
401                matches!(target_dtype, DType::U8 | DType::Bool(BoolStore::U8))
402            }
403            DType::Bool(BoolStore::U32) => {
404                matches!(target_dtype, DType::U32 | DType::Bool(BoolStore::U32))
405            }
406            dtype => dtype == target_dtype,
407        }
408    }
409
410    /// Returns the rank (the number of dimensions).
411    pub fn rank(&self) -> usize {
412        self.shape.len()
413    }
414
415    /// Returns the total number of elements of the tensor data.
416    pub fn num_elements(&self) -> usize {
417        numel(&self.shape)
418    }
419
420    /// Populates the data with random values.
421    pub fn random<E: Element, R: Rng, S: Into<Shape>>(
422        shape: S,
423        distribution: Distribution,
424        rng: &mut R,
425    ) -> Self {
426        let shape = shape.into();
427        let num_elements = numel(&shape);
428        let mut data = Vec::with_capacity(num_elements);
429
430        for _ in 0..num_elements {
431            data.push(E::random(distribution, rng));
432        }
433
434        TensorData::new(data, shape)
435    }
436
437    /// Populates the data with random values of the requested dtype.
438    ///
439    /// Boolean values are sampled as `bool` and encoded using the requested storage type.
440    ///
441    /// # Panics
442    ///
443    /// Panics if `dtype` is [`DType::QFloat`] or the distribution is invalid for the element type.
444    pub fn random_dtype<R: Rng, S: Into<Shape>>(
445        shape: S,
446        distribution: Distribution,
447        rng: &mut R,
448        dtype: DType,
449    ) -> Self {
450        match dtype {
451            DType::F64 => Self::random::<f64, _, _>(shape, distribution, rng),
452            DType::F32 => Self::random::<f32, _, _>(shape, distribution, rng),
453            DType::Flex32 => Self::random::<flex32, _, _>(shape, distribution, rng),
454            DType::F16 => Self::random::<f16, _, _>(shape, distribution, rng),
455            DType::BF16 => Self::random::<bf16, _, _>(shape, distribution, rng),
456            DType::I64 => Self::random::<i64, _, _>(shape, distribution, rng),
457            DType::I32 => Self::random::<i32, _, _>(shape, distribution, rng),
458            DType::I16 => Self::random::<i16, _, _>(shape, distribution, rng),
459            DType::I8 => Self::random::<i8, _, _>(shape, distribution, rng),
460            DType::U64 => Self::random::<u64, _, _>(shape, distribution, rng),
461            DType::U32 => Self::random::<u32, _, _>(shape, distribution, rng),
462            DType::U16 => Self::random::<u16, _, _>(shape, distribution, rng),
463            DType::U8 => Self::random::<u8, _, _>(shape, distribution, rng),
464            DType::Bool(_) => {
465                Self::random::<bool, _, _>(shape, distribution, rng).convert_dtype(dtype)
466            }
467            DType::QFloat(_) => {
468                panic!("Cannot create random tensor data with quantized dtype {dtype:?}")
469            }
470        }
471    }
472
473    /// Populates the data with zeros.
474    pub fn zeros<E: Element, S: Into<Shape>>(shape: S) -> TensorData {
475        Self::full(shape, 0.elem::<E>())
476    }
477
478    /// Populates the data with zeros of the requested dtype.
479    ///
480    /// # Panics
481    ///
482    /// Panics if `dtype` is [`DType::QFloat`].
483    pub fn zeros_dtype<S: Into<Shape>>(shape: S, dtype: DType) -> Self {
484        Self::full_dtype(shape, 0, dtype)
485    }
486
487    /// Populates the data with ones.
488    pub fn ones<E: Element, S: Into<Shape>>(shape: S) -> TensorData {
489        Self::full(shape, 1.elem::<E>())
490    }
491
492    /// Populates the data with ones of the requested dtype.
493    ///
494    /// # Panics
495    ///
496    /// Panics if `dtype` is [`DType::QFloat`].
497    pub fn ones_dtype<S: Into<Shape>>(shape: S, dtype: DType) -> Self {
498        Self::full_dtype(shape, 1, dtype)
499    }
500
501    /// Populates the data with the given value
502    pub fn full<E: Element, S: Into<Shape>>(shape: S, fill_value: E) -> TensorData {
503        let shape = shape.into();
504        let num_elements = numel(&shape);
505        let mut data = Vec::<E>::with_capacity(num_elements);
506        for _ in 0..num_elements {
507            data.push(fill_value)
508        }
509
510        TensorData::new(data, shape)
511    }
512
513    /// Populates the data with the given value of the requested dtype.
514    ///
515    /// # Panics
516    ///
517    /// Panics if `dtype` is [`DType::QFloat`].
518    pub fn full_dtype<E: Into<Scalar>, S: Into<Shape>>(
519        shape: S,
520        fill_value: E,
521        dtype: DType,
522    ) -> TensorData {
523        let fill_value = fill_value.into();
524        match dtype {
525            DType::F64 => Self::full::<f64, _>(shape, fill_value.elem()),
526            DType::F32 => Self::full::<f32, _>(shape, fill_value.elem()),
527            DType::Flex32 => Self::full::<flex32, _>(shape, fill_value.elem()),
528            DType::F16 => Self::full::<f16, _>(shape, fill_value.elem()),
529            DType::BF16 => Self::full::<bf16, _>(shape, fill_value.elem()),
530            DType::I64 => Self::full::<i64, _>(shape, fill_value.elem()),
531            DType::I32 => Self::full::<i32, _>(shape, fill_value.elem()),
532            DType::I16 => Self::full::<i16, _>(shape, fill_value.elem()),
533            DType::I8 => Self::full::<i8, _>(shape, fill_value.elem()),
534            DType::U64 => Self::full::<u64, _>(shape, fill_value.elem()),
535            DType::U32 => Self::full::<u32, _>(shape, fill_value.elem()),
536            DType::U16 => Self::full::<u16, _>(shape, fill_value.elem()),
537            DType::U8 => Self::full::<u8, _>(shape, fill_value.elem()),
538            DType::Bool(BoolStore::Native) => Self::full::<bool, _>(shape, fill_value.elem()),
539            DType::Bool(BoolStore::U8) => {
540                Self::full::<u8, _>(shape, fill_value.elem()).into_bool_u8()
541            }
542            DType::Bool(BoolStore::U32) => {
543                Self::full::<u32, _>(shape, fill_value.elem()).into_bool_u32()
544            }
545            DType::QFloat(_) => {
546                panic!("Cannot create full tensor data with quantized dtype {dtype:?}")
547            }
548        }
549    }
550
551    // Unchecked, used to overwrite the dtype
552    pub(super) fn into_bool_u8(mut self) -> Self {
553        self.dtype = DType::Bool(BoolStore::U8);
554        self
555    }
556
557    // Unchecked, used to overwrite the dtype
558    pub(super) fn into_bool_u32(mut self) -> Self {
559        self.dtype = DType::Bool(BoolStore::U32);
560        self
561    }
562
563    /// Returns the data as a slice of bytes.
564    pub fn as_bytes(&self) -> &[u8] {
565        &self.bytes
566    }
567
568    /// Returns the bytes representation of the data.
569    pub fn into_bytes(self) -> Bytes {
570        self.bytes
571    }
572}
573
574/// Resets every item whose byte length no longer matches the length it had when the guard was
575/// created, including when dropped during unwinding.
576struct LengthGuard<'a> {
577    data: Vec<&'a mut TensorData>,
578    lens: Vec<usize>,
579}
580
581impl<'a> LengthGuard<'a> {
582    fn new(data: Vec<&'a mut TensorData>) -> Self {
583        let lens = data.iter().map(|data| data.bytes.len()).collect();
584        Self { data, lens }
585    }
586
587    /// Replaces changed items with empty data and returns whether any changed.
588    fn reset_changed(&mut self) -> bool {
589        let mut changed = false;
590        for (data, len) in self.data.iter_mut().zip(&self.lens) {
591            if data.bytes.len() != *len {
592                **data = TensorData::new(Vec::<u8>::new(), [0]);
593                changed = true;
594            }
595        }
596        changed
597    }
598}
599
600impl Drop for LengthGuard<'_> {
601    fn drop(&mut self) {
602        self.reset_changed();
603    }
604}
605
606/// Number of elements described by `shape`, or `None` if the product overflows `usize`.
607fn checked_numel(shape: &[usize]) -> Option<usize> {
608    shape
609        .iter()
610        .try_fold(1usize, |numel, dim| numel.checked_mul(*dim))
611}
612
613/// Number of elements described by `shape`.
614///
615/// Panics if the product overflows `usize`, since a wrapped count could otherwise match a
616/// shorter buffer.
617fn numel(shape: &[usize]) -> usize {
618    match checked_numel(shape) {
619        Some(numel) => numel,
620        None => panic!("Shape {shape:?} has more elements than fit in usize"),
621    }
622}
623
624impl<E: Element, const A: usize> From<[E; A]> for TensorData {
625    fn from(elems: [E; A]) -> Self {
626        TensorData::new(elems.to_vec(), [A])
627    }
628}
629
630impl<const A: usize> From<[usize; A]> for TensorData {
631    fn from(elems: [usize; A]) -> Self {
632        TensorData::new(elems.iter().map(|&e| e as i64).collect(), [A])
633    }
634}
635
636impl From<&[usize]> for TensorData {
637    fn from(elems: &[usize]) -> Self {
638        let mut data = Vec::with_capacity(elems.len());
639        for elem in elems.iter() {
640            data.push(*elem as i64);
641        }
642
643        TensorData::new(data, [elems.len()])
644    }
645}
646
647impl<E: Element> From<&[E]> for TensorData {
648    fn from(elems: &[E]) -> Self {
649        let mut data = Vec::with_capacity(elems.len());
650        for elem in elems.iter() {
651            data.push(*elem);
652        }
653
654        TensorData::new(data, [elems.len()])
655    }
656}
657
658impl<E: Element, const A: usize, const B: usize> From<[[E; B]; A]> for TensorData {
659    fn from(elems: [[E; B]; A]) -> Self {
660        let mut data = Vec::with_capacity(A * B);
661        for elem in elems.into_iter().take(A) {
662            for elem in elem.into_iter().take(B) {
663                data.push(elem);
664            }
665        }
666
667        TensorData::new(data, [A, B])
668    }
669}
670
671impl<E: Element, const A: usize, const B: usize, const C: usize> From<[[[E; C]; B]; A]>
672    for TensorData
673{
674    fn from(elems: [[[E; C]; B]; A]) -> Self {
675        let mut data = Vec::with_capacity(A * B * C);
676
677        for elem in elems.into_iter().take(A) {
678            for elem in elem.into_iter().take(B) {
679                for elem in elem.into_iter().take(C) {
680                    data.push(elem);
681                }
682            }
683        }
684
685        TensorData::new(data, [A, B, C])
686    }
687}
688
689impl<E: Element, const A: usize, const B: usize, const C: usize, const D: usize>
690    From<[[[[E; D]; C]; B]; A]> for TensorData
691{
692    fn from(elems: [[[[E; D]; C]; B]; A]) -> Self {
693        let mut data = Vec::with_capacity(A * B * C * D);
694
695        for elem in elems.into_iter().take(A) {
696            for elem in elem.into_iter().take(B) {
697                for elem in elem.into_iter().take(C) {
698                    for elem in elem.into_iter().take(D) {
699                        data.push(elem);
700                    }
701                }
702            }
703        }
704
705        TensorData::new(data, [A, B, C, D])
706    }
707}
708
709impl<Elem: Element, const A: usize, const B: usize, const C: usize, const D: usize, const E: usize>
710    From<[[[[[Elem; E]; D]; C]; B]; A]> for TensorData
711{
712    fn from(elems: [[[[[Elem; E]; D]; C]; B]; A]) -> Self {
713        let mut data = Vec::with_capacity(A * B * C * D * E);
714
715        for elem in elems.into_iter().take(A) {
716            for elem in elem.into_iter().take(B) {
717                for elem in elem.into_iter().take(C) {
718                    for elem in elem.into_iter().take(D) {
719                        for elem in elem.into_iter().take(E) {
720                            data.push(elem);
721                        }
722                    }
723                }
724            }
725        }
726
727        TensorData::new(data, [A, B, C, D, E])
728    }
729}
730impl core::fmt::Display for TensorData {
731    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
732        let fmt = match self.dtype {
733            DType::F64 => format!("{:?}", self.as_slice::<f64>().unwrap()),
734            DType::F32 | DType::Flex32 => format!("{:?}", self.as_slice::<f32>().unwrap()),
735            DType::F16 => format!("{:?}", self.as_slice::<f16>().unwrap()),
736            DType::BF16 => format!("{:?}", self.as_slice::<bf16>().unwrap()),
737            DType::I64 => format!("{:?}", self.as_slice::<i64>().unwrap()),
738            DType::I32 => format!("{:?}", self.as_slice::<i32>().unwrap()),
739            DType::I16 => format!("{:?}", self.as_slice::<i16>().unwrap()),
740            DType::I8 => format!("{:?}", self.as_slice::<i8>().unwrap()),
741            DType::U64 => format!("{:?}", self.as_slice::<u64>().unwrap()),
742            DType::U32 => format!("{:?}", self.as_slice::<u32>().unwrap()),
743            DType::U16 => format!("{:?}", self.as_slice::<u16>().unwrap()),
744            DType::U8 => format!("{:?}", self.as_slice::<u8>().unwrap()),
745            DType::Bool(BoolStore::Native) => format!("{:?}", self.as_slice::<bool>().unwrap()),
746            DType::Bool(BoolStore::U8) => format!("{:?}", self.as_slice::<u8>().unwrap()),
747            DType::Bool(BoolStore::U32) => format!("{:?}", self.as_slice::<u32>().unwrap()),
748            DType::QFloat(scheme) => match scheme {
749                QuantScheme {
750                    mode: QuantMode::Symmetric,
751                    value:
752                        QuantValue::Q8F
753                        | QuantValue::Q8S
754                        // Display sub-byte values as i8
755                        | QuantValue::Q4F
756                        | QuantValue::Q4S
757                        | QuantValue::Q2F
758                        | QuantValue::Q2S,
759                    ..
760                } => {
761                    format!("{:?} {scheme:?}", self.iter::<i8>().collect::<Vec<_>>())
762                },
763                QuantScheme {
764                        mode: QuantMode::Symmetric,
765                        value:
766                            QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1,
767                        ..
768                    } => {
769                        format!("<float-quantized> {scheme:?}")
770                    }
771                QuantScheme {
772                    mode: QuantMode::Lookup,
773                    ..
774                } => {
775                    format!("<lookup-quantized> {scheme:?}")
776                }
777            },
778        };
779        f.write_str(fmt.as_str())
780    }
781}
782
783#[cfg(test)]
784mod tests {
785    use super::*;
786    use crate::*;
787    use ::rand::{
788        SeedableRng,
789        rngs::{StdRng, SysRng},
790    };
791    use alloc::string::ToString;
792    use alloc::vec;
793    use core::mem::{MaybeUninit, align_of, size_of};
794
795    #[test]
796    #[should_panic(expected = "Cannot create full tensor data with quantized dtype")]
797    fn full_dtype_rejects_quantized() {
798        TensorData::full_dtype([2, 3], 0, DType::QFloat(Default::default()));
799    }
800
801    #[test]
802    #[should_panic(expected = "Cannot create random tensor data with quantized dtype")]
803    fn random_dtype_rejects_quantized() {
804        TensorData::random_dtype(
805            [2, 3],
806            Distribution::Default,
807            &mut StdRng::seed_from_u64(42),
808            DType::QFloat(Default::default()),
809        );
810    }
811
812    #[test]
813    fn should_have_rank() {
814        let shape = [3, 5, 6];
815        let data = TensorData::random::<f32, _, _>(
816            shape,
817            Distribution::Default,
818            &mut StdRng::try_from_rng(&mut SysRng).unwrap(),
819        );
820
821        assert_eq!(data.rank(), 3);
822    }
823
824    #[test]
825    fn into_vec_should_yield_same_value_as_iter() {
826        let shape = [3, 5, 6];
827        let data = TensorData::random::<f32, _, _>(
828            shape,
829            Distribution::Default,
830            &mut StdRng::try_from_rng(&mut SysRng).unwrap(),
831        );
832
833        let expected = data.iter::<f32>().collect::<Vec<f32>>();
834        let actual = data.try_into_vec::<f32>().unwrap();
835
836        assert_eq!(expected, actual);
837    }
838
839    #[test]
840    #[should_panic]
841    fn into_vec_should_assert_wrong_dtype() {
842        let shape = [3, 5, 6];
843        let data = TensorData::random::<f32, _, _>(
844            shape,
845            Distribution::Default,
846            &mut StdRng::try_from_rng(&mut SysRng).unwrap(),
847        );
848
849        data.try_into_vec::<i32>().unwrap();
850    }
851
852    #[test]
853    fn should_have_right_num_elements() {
854        let shape = [3, 5, 6];
855        let num_elements: usize = shape.iter().product();
856        let data = TensorData::random::<f32, _, _>(
857            shape,
858            Distribution::Default,
859            &mut StdRng::try_from_rng(&mut SysRng).unwrap(),
860        );
861
862        assert_eq!(num_elements, data.bytes.len() / 4); // f32 stored as u8s
863        assert_eq!(num_elements, data.as_slice::<f32>().unwrap().len());
864    }
865
866    #[test]
867    fn should_have_right_shape() {
868        let data = TensorData::from([[3.0, 5.0, 6.0]]);
869        assert_eq!(data.shape, shape![1, 3]);
870
871        let data = TensorData::from([[4.0, 5.0, 8.0], [3.0, 5.0, 6.0]]);
872        assert_eq!(data.shape, shape![2, 3]);
873
874        let data = TensorData::from([3.0, 5.0, 6.0]);
875        assert_eq!(data.shape, shape![3]);
876    }
877
878    #[test]
879    fn should_convert_bytes_correctly() {
880        let mut vector: Vec<f32> = Vec::with_capacity(5);
881        vector.push(2.0);
882        vector.push(3.0);
883        let data1 = TensorData::new(vector, vec![2]);
884
885        let factor = size_of::<f32>() / size_of::<u8>();
886        assert_eq!(data1.bytes.len(), 2 * factor);
887        assert_eq!(data1.bytes.capacity(), 5 * factor);
888    }
889
890    #[test]
891    fn should_convert_bytes_correctly_inplace() {
892        fn test_precision<E: Element>() {
893            let data = TensorData::new((0..32).collect(), [32]);
894            let self1 = data.clone().convert::<E>();
895            for (i, val) in self1.try_into_vec::<E>().unwrap().into_iter().enumerate() {
896                assert_eq!(i as u32, val.elem::<u32>())
897            }
898        }
899        test_precision::<f32>();
900        test_precision::<f16>();
901        test_precision::<i64>();
902        test_precision::<i32>();
903    }
904
905    #[test]
906    fn should_preserve_flex32_dtype_when_casting_from_f32() {
907        let data = TensorData::from([-1.5f32, 0.0, 2.25]);
908        let expected = TensorData::from([-1.5, 0.0, 2.25].map(flex32::from_f32));
909
910        let actual = data.clone().try_cast(DType::Flex32).unwrap();
911        assert_eq!(actual, expected);
912        assert_eq!(data.try_cast_as::<flex32>().unwrap(), expected);
913        assert_eq!(
914            actual.try_cast(DType::F32).unwrap(),
915            TensorData::from([-1.5f32, 0.0, 2.25])
916        );
917    }
918
919    #[test]
920    fn should_preserve_flex32_dtype_when_casting_from_f64() {
921        let data = TensorData::from([-1.5f64, 0.0, 2.25]);
922        let expected = TensorData::from([-1.5, 0.0, 2.25].map(flex32::from_f32));
923
924        let actual = data.clone().try_cast(DType::Flex32).unwrap();
925        assert_eq!(actual, expected);
926        assert_eq!(data.try_cast_as::<flex32>().unwrap(), expected);
927        assert_eq!(
928            actual.try_cast(DType::F32).unwrap(),
929            TensorData::from([-1.5f32, 0.0, 2.25])
930        );
931    }
932
933    #[test]
934    fn should_preserve_flex32_dtype_when_constructing_full_data() {
935        let actual = TensorData::full_dtype([2, 2], 1.5f32, DType::Flex32);
936        let expected = TensorData::full::<flex32, _>([2, 2], 1.5f32.elem());
937
938        assert_eq!(actual, expected);
939        assert_eq!(actual.dtype(), DType::Flex32);
940    }
941
942    #[test]
943    fn should_convert_negative_values_to_bool_store() {
944        for store in [BoolStore::U8, BoolStore::U32, BoolStore::Native] {
945            let data = TensorData::from([-1i32, 0, 1, -12]).convert_dtype(DType::Bool(store));
946            assert_eq!(data.dtype, DType::Bool(store));
947            assert_eq!(
948                data.iter::<bool>().collect::<Vec<_>>(),
949                [true, false, true, true]
950            );
951
952            let data = TensorData::from([-1.5f32, 0.0, 0.5]).convert_dtype(DType::Bool(store));
953            assert_eq!(data.iter::<bool>().collect::<Vec<_>>(), [true, false, true]);
954        }
955    }
956
957    fn check_dtype_constructors<E: Element + PartialEq, S: Element>(dtype: DType) {
958        let distribution = Distribution::Uniform(0.0, 8.0);
959        let random =
960            TensorData::random_dtype([2, 3], distribution, &mut StdRng::seed_from_u64(42), dtype);
961        let expected =
962            TensorData::random::<S, _, _>([2, 3], distribution, &mut StdRng::seed_from_u64(42));
963        for (data, values) in [
964            (
965                TensorData::zeros_dtype([2, 3], dtype),
966                vec![0.elem::<E>(); 6],
967            ),
968            (
969                TensorData::ones_dtype([2, 3], dtype),
970                vec![1.elem::<E>(); 6],
971            ),
972            (random, expected.iter::<E>().collect()),
973        ] {
974            assert_eq!(data.dtype(), dtype);
975            assert_eq!(data.shape(), &Shape::new([2, 3]));
976            assert_eq!(data.try_into_vec::<E>().unwrap(), values);
977        }
978    }
979
980    macro_rules! test_dtypes {
981    ($test_name:ident, $($dtype:ty),*) => {
982        $(
983            paste::paste! {
984                #[test]
985                fn [<$test_name _ $dtype:snake>]() {
986                    check_dtype_constructors::<$dtype, $dtype>(<$dtype>::dtype());
987                    let full_dtype = TensorData::full_dtype([2, 16], 4, <$dtype>::dtype());
988                    let full = TensorData::full::<$dtype, _>([2, 16], 4.elem());
989                    assert_eq!(full_dtype, full);
990                }
991            }
992        )*
993    };
994}
995
996    test_dtypes!(
997        should_create_with_dtype,
998        bool,
999        i8,
1000        i16,
1001        i32,
1002        i64,
1003        u8,
1004        u16,
1005        u32,
1006        u64,
1007        f16,
1008        bf16,
1009        f32,
1010        f64,
1011        flex32
1012    );
1013
1014    #[test]
1015    fn should_create_with_dtype_bool_u8() {
1016        check_dtype_constructors::<u8, bool>(DType::Bool(BoolStore::U8));
1017    }
1018
1019    #[test]
1020    fn should_create_with_dtype_bool_u32() {
1021        check_dtype_constructors::<u32, bool>(DType::Bool(BoolStore::U32));
1022    }
1023
1024    #[test]
1025    fn should_serialize_deserialize_tensor_data() {
1026        let data = TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3]);
1027        assert_eq!(
1028            data.as_bytes(),
1029            [
1030                0, 0, 128, 63, 0, 0, 0, 64, 0, 0, 64, 64, 0, 0, 128, 64, 0, 0, 160, 64, 0, 0, 192,
1031                64
1032            ]
1033        );
1034        let serialized = serde_json::to_string(&data).unwrap();
1035        let deserialized: TensorData = serde_json::from_str(&serialized).unwrap();
1036        assert_eq!(data, deserialized);
1037    }
1038
1039    #[test]
1040    fn should_deserialize_tensor_data_with_shape_inner() {
1041        // TensorData `shape` was previously a Vec<usize>.
1042        let serialized = r#"{
1043        "bytes": [0, 0, 128, 63, 0, 0, 0, 64, 0, 0, 64, 64, 0, 0, 128, 64, 0, 0, 160, 64, 0, 0, 192, 64],
1044        "shape": [2, 3],
1045        "dtype": "F32"
1046    }"#;
1047
1048        let data: TensorData = serde_json::from_str(serialized).unwrap();
1049        assert_eq!(data.shape, shape![2, 3]);
1050        assert_eq!(
1051            data.as_slice::<f32>().unwrap(),
1052            &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
1053        );
1054    }
1055
1056    #[test]
1057    fn should_not_deserialize_tensor_data_with_shape_larger_than_bytes() {
1058        // A crafted record whose shape claims 1000 f32 over a 4-byte payload must be rejected,
1059        // otherwise the resulting tensor reads and writes out of bounds.
1060        let serialized = r#"{"bytes": [0, 0, 128, 63], "shape": [1000], "dtype": "F32"}"#;
1061
1062        let err = serde_json::from_str::<TensorData>(serialized).unwrap_err();
1063        assert!(
1064            err.to_string().contains("is invalid for input of size"),
1065            "unexpected error: {err}"
1066        );
1067    }
1068
1069    #[test]
1070    fn should_not_deserialize_tensor_data_with_overflowing_shape() {
1071        // The dimensions multiply to 2 with a wrapped product, matching the payload, so an
1072        // unchecked product would accept a shape that claims 2^63 + 1 rows.
1073        let serialized = r#"{"bytes": [0, 0, 128, 63, 0, 0, 0, 64], "shape": [9223372036854775809, 2], "dtype": "F32"}"#;
1074
1075        let err = serde_json::from_str::<TensorData>(serialized).unwrap_err();
1076        assert!(
1077            err.to_string().contains("is invalid for input of size"),
1078            "unexpected error: {err}"
1079        );
1080    }
1081
1082    #[test]
1083    fn should_serialize_shape_as_flat_array() {
1084        // Ensure the new Shape serializes identically to how Vec<usize> used to,
1085        // i.e. as a flat JSON array, not as an object like `{"dims": [2, 3]}`.
1086        let data = TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3]);
1087        let serialized = serde_json::to_string(&data).unwrap();
1088        let json: serde_json::Value = serde_json::from_str(&serialized).unwrap();
1089        assert_eq!(json["shape"], serde_json::json!([2, 3]));
1090    }
1091
1092    #[test]
1093    fn test_tensor_data_try_view_dtype_mismatch() {
1094        let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]);
1095        let dtype = data.dtype;
1096
1097        assert_eq!(
1098            data.try_view::<i32>().unwrap_err(),
1099            DataError::DTypeMismatch {
1100                expected: <i32 as Element>::dtype(),
1101                actual: dtype,
1102            }
1103        );
1104    }
1105
1106    #[test]
1107    #[should_panic(expected = "Expected data type")]
1108    fn test_tensor_data_expect_view_dtype_mismatch() {
1109        let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]);
1110        let _view = data.view::<i32>();
1111    }
1112
1113    #[test]
1114    fn test_tensor_data_try_mut_view_dtype_mismatch() {
1115        let mut data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]);
1116
1117        let result = data.try_mut_view::<i32>();
1118        assert_eq!(
1119            result.unwrap_err(),
1120            DataError::DTypeMismatch {
1121                actual: data.dtype,
1122                expected: <i32 as Element>::dtype(),
1123            }
1124        );
1125    }
1126
1127    #[test]
1128    fn try_view_validates_storage() {
1129        let invalid_representation =
1130            TensorData::from_bytes_unchecked(Bytes::from_bytes_vec(vec![0; 3]), [1], DType::F32);
1131        assert!(matches!(
1132            invalid_representation.try_view::<f32>(),
1133            Err(DataError::InvalidRepresentation(_))
1134        ));
1135
1136        let invalid_count =
1137            TensorData::from_bytes_unchecked(Bytes::from_bytes_vec(vec![0; 8]), [1], DType::F32);
1138        assert!(matches!(
1139            invalid_count.try_view::<f32>(),
1140            Err(DataError::ElementCountMismatch {
1141                expected: 1,
1142                actual: 2
1143            })
1144        ));
1145
1146        let invalid_bool = TensorData::from_bytes_vec(vec![2], [1], DType::Bool(BoolStore::Native));
1147        assert!(matches!(
1148            invalid_bool.try_view::<bool>(),
1149            Err(DataError::InvalidRepresentation(_))
1150        ));
1151    }
1152
1153    #[test]
1154    fn try_from_bytes_rejects_invalid_length() {
1155        let err = TensorData::try_from_bytes_vec(vec![0; 3], [1], DType::F32).unwrap_err();
1156        assert_eq!(
1157            err,
1158            DataError::InvalidByteLength {
1159                shape: Shape::from([1]),
1160                dtype: DType::F32,
1161                len: 3,
1162            }
1163        );
1164
1165        assert!(TensorData::try_from_bytes_vec(vec![0; 8], [1], DType::F32).is_err());
1166        assert!(TensorData::try_from_bytes_vec(vec![], [1], DType::F32).is_err());
1167        assert!(TensorData::try_from_bytes_vec(vec![], [0usize; 0], DType::F32).is_err());
1168        // Wrapped product of the dimensions matches the payload.
1169        assert!(
1170            TensorData::try_from_bytes_vec(vec![0; 8], [usize::MAX / 2 + 2, 2], DType::F32)
1171                .is_err()
1172        );
1173    }
1174
1175    #[test]
1176    #[should_panic(expected = "more elements than fit in usize")]
1177    fn new_rejects_overflowing_shape() {
1178        // The wrapped product (2) matches the element count.
1179        let _ = TensorData::new(vec![0.0f32; 2], [usize::MAX / 2 + 2, 2]);
1180    }
1181
1182    #[test]
1183    #[should_panic(expected = "more elements than fit in usize")]
1184    fn zeros_rejects_overflowing_shape() {
1185        let _ = TensorData::zeros::<f32, _>([usize::MAX / 2 + 2, 2]);
1186    }
1187
1188    #[test]
1189    fn try_from_bytes_accepts_valid_length() {
1190        let data = TensorData::try_from_bytes_vec(vec![0; 8], [2], DType::F32).unwrap();
1191        assert_eq!(data.shape(), &Shape::from([2]));
1192        assert_eq!(data.dtype(), DType::F32);
1193        assert_eq!(data.bytes().len(), 8);
1194
1195        let empty = TensorData::try_from_bytes_vec(vec![], [0, 3], DType::F32).unwrap();
1196        assert_eq!(empty.num_elements(), 0);
1197
1198        let scalar = TensorData::try_from_bytes_vec(vec![0; 4], [0usize; 0], DType::F32).unwrap();
1199        assert_eq!(scalar.rank(), 0);
1200    }
1201
1202    #[test]
1203    #[should_panic(expected = "is invalid for input of size 3 bytes")]
1204    fn from_bytes_panics_on_invalid_length() {
1205        let _ = TensorData::from_bytes_vec(vec![0; 3], [1], DType::F32);
1206    }
1207
1208    #[test]
1209    fn try_from_bytes_accepts_quantized_layout() {
1210        // Quantized storage holds the values followed by the scales, so its length isn't
1211        // `numel * dtype.size()`.
1212        let data = TensorData::quantized(vec![0i8; 4], [4], QuantScheme::default(), &[1.0], None);
1213        let (bytes, shape, dtype) = data.clone().into_parts();
1214        assert_eq!(
1215            TensorData::try_from_bytes(bytes, shape, dtype).unwrap(),
1216            data
1217        );
1218
1219        let serialized = serde_json::to_string(&data).unwrap();
1220        let deserialized: TensorData = serde_json::from_str(&serialized).unwrap();
1221        assert_eq!(deserialized, data);
1222    }
1223
1224    #[test]
1225    fn packed_quantized_scalar_roundtrips() {
1226        for store in [QuantStore::PackedU32(0), QuantStore::PackedNative(0)] {
1227            let scheme = QuantScheme::default()
1228                .with_store(store)
1229                .with_value(QuantValue::E2M1);
1230            let data = TensorData::quantized(vec![1i8], [0usize; 0], scheme, &[0.5], None);
1231
1232            let serialized = serde_json::to_string(&data).unwrap();
1233            let deserialized: TensorData = serde_json::from_str(&serialized).unwrap();
1234            assert_eq!(deserialized, data);
1235        }
1236    }
1237
1238    #[test]
1239    fn try_from_bytes_rejects_invalid_quantized_lengths_and_schemes() {
1240        let scheme = QuantScheme::default();
1241        let valid = TensorData::quantized(vec![0i8; 3], [3], scheme, &[1.0], None);
1242        let (bytes, shape, dtype) = valid.into_parts();
1243        assert!(TensorData::try_from_bytes(bytes.clone(), shape.clone(), dtype).is_ok());
1244
1245        for len in [bytes.len() - 1, bytes.len() + 1] {
1246            assert!(matches!(
1247                TensorData::try_from_bytes_vec(vec![0; len], shape.clone(), dtype),
1248                Err(DataError::InvalidByteLength { .. })
1249            ));
1250        }
1251
1252        let invalid_packed_dim = scheme.with_store(QuantStore::PackedU32(1));
1253        for shape in [Shape::new([1]), Shape::new([0usize; 0])] {
1254            assert!(matches!(
1255                TensorData::try_from_bytes_vec(
1256                    vec![0; 4],
1257                    shape,
1258                    DType::QFloat(invalid_packed_dim)
1259                ),
1260                Err(DataError::InvalidQuantizedLayout { .. })
1261            ));
1262        }
1263
1264        let zero_block = scheme.per_block([0], crate::ScaleDtype::F32);
1265        assert!(matches!(
1266            TensorData::try_from_bytes_vec(vec![0; 4], [1], DType::QFloat(zero_block)),
1267            Err(DataError::InvalidQuantizedLayout { .. })
1268        ));
1269
1270        assert!(matches!(
1271            TensorData::try_from_bytes_vec(vec![], [usize::MAX, 2], DType::QFloat(scheme)),
1272            Err(DataError::InvalidQuantizedLayout { .. })
1273        ));
1274    }
1275
1276    #[test]
1277    fn deserialization_rejects_invalid_quantized_length() {
1278        let invalid = TensorData::from_bytes_unchecked(
1279            Bytes::from_bytes_vec(vec![0; 3]),
1280            [4],
1281            DType::QFloat(QuantScheme::default().with_store(QuantStore::Native)),
1282        );
1283        let serialized = serde_json::to_string(&invalid).unwrap();
1284
1285        assert!(serde_json::from_str::<TensorData>(&serialized).is_err());
1286    }
1287
1288    #[test]
1289    fn try_from_bytes_accepts_packed_quantized_layout_per_row() {
1290        let scheme = QuantScheme::default()
1291            .with_value(QuantValue::Q4S)
1292            .with_store(QuantStore::PackedU32(0));
1293        let data = TensorData::quantized(vec![0i8; 9], [3, 3], scheme, &[1.0], None);
1294        let (bytes, shape, dtype) = data.into_parts();
1295
1296        // Three rows each occupy one u32 word, followed by the per-tensor f32 scale.
1297        assert_eq!(bytes.len(), 3 * core::mem::size_of::<u32>() + 4);
1298        assert!(TensorData::try_from_bytes(bytes.clone(), shape.clone(), dtype).is_ok());
1299        for len in [bytes.len() - 1, bytes.len() + 1] {
1300            assert!(matches!(
1301                TensorData::try_from_bytes_vec(vec![0; len], shape.clone(), dtype),
1302                Err(DataError::InvalidByteLength { .. })
1303            ));
1304        }
1305    }
1306
1307    #[test]
1308    fn try_from_bytes_accepts_zero_sized_quantized_tensors() {
1309        let per_tensor =
1310            TensorData::quantized(Vec::<i8>::new(), [0], QuantScheme::default(), &[1.0], None);
1311        let (bytes, shape, dtype) = per_tensor.into_parts();
1312        assert_eq!(bytes.len(), 4);
1313        assert!(TensorData::try_from_bytes(bytes, shape, dtype).is_ok());
1314
1315        let per_block = QuantScheme::default()
1316            .with_store(QuantStore::Native)
1317            .per_block([2], crate::ScaleDtype::F32);
1318        let empty = TensorData::quantized(Vec::<i8>::new(), [0], per_block, &[], None);
1319        let (bytes, shape, dtype) = empty.into_parts();
1320        assert!(bytes.is_empty());
1321        assert!(TensorData::try_from_bytes(bytes, shape, dtype).is_ok());
1322    }
1323
1324    #[test]
1325    fn with_bytes_mut_allows_length_preserving_writes() {
1326        let mut a = TensorData::from([1.0f32, 2.0]);
1327        let mut b = TensorData::from([3i32]);
1328
1329        TensorData::with_bytes_mut([&mut a, &mut b], |bytes| {
1330            assert_eq!(bytes.len(), 2);
1331            for bytes in bytes {
1332                bytes.fill(0);
1333            }
1334        });
1335
1336        assert_eq!(a, TensorData::from([0.0f32, 0.0]));
1337        assert_eq!(b, TensorData::from([0i32]));
1338    }
1339
1340    #[test]
1341    #[should_panic(expected = "length must not change")]
1342    fn with_bytes_mut_rejects_length_change() {
1343        let mut data = TensorData::from([1.0f32]);
1344        TensorData::with_bytes_mut([&mut data], |mut bytes| {
1345            bytes[0].extend_from_byte_slice(&[0; 4]);
1346        });
1347    }
1348
1349    #[cfg(feature = "std")]
1350    #[test]
1351    fn with_bytes_mut_leaves_no_invalid_data_after_unwind() {
1352        use std::panic::{AssertUnwindSafe, catch_unwind};
1353
1354        let empty = TensorData::new(Vec::<u8>::new(), [0]);
1355
1356        // Length change detected after `f` returns.
1357        let mut changed = TensorData::from([1.0f32]);
1358        let mut unchanged = TensorData::from([2.0f32]);
1359        let result = catch_unwind(AssertUnwindSafe(|| {
1360            TensorData::with_bytes_mut([&mut changed, &mut unchanged], |mut bytes| {
1361                bytes[0].extend_from_byte_slice(&[0; 4]);
1362            })
1363        }));
1364        assert!(result.is_err());
1365        assert_eq!(changed, empty);
1366        assert_eq!(unchanged, TensorData::from([2.0f32]));
1367
1368        // `f` panics after changing a length.
1369        let mut data = TensorData::from([1.0f32]);
1370        let result = catch_unwind(AssertUnwindSafe(|| {
1371            TensorData::with_bytes_mut([&mut data], |mut bytes| {
1372                bytes[0].extend_from_byte_slice(&[0; 4]);
1373                panic!("callback failed");
1374            })
1375        }));
1376        assert!(result.is_err());
1377        assert_eq!(data, empty);
1378    }
1379
1380    #[test]
1381    fn try_view_propagates_storage_access_failure() {
1382        use alloc::boxed::Box;
1383
1384        #[derive(Debug)]
1385        struct FailingController;
1386
1387        impl AllocationController for FailingController {
1388            fn alloc_align(&self) -> usize {
1389                align_of::<f32>()
1390            }
1391
1392            fn property(&self) -> AllocationProperty {
1393                AllocationProperty::Other
1394            }
1395
1396            fn capacity(&self) -> usize {
1397                size_of::<f32>()
1398            }
1399
1400            fn memory(&self, _policy: AccessPolicy) -> Result<&[MaybeUninit<u8>], AccessError> {
1401                Err(AccessError::Read("test read failure".into()))
1402            }
1403
1404            unsafe fn memory_mut(
1405                &mut self,
1406                _policy: AccessPolicy,
1407            ) -> Result<&mut [MaybeUninit<u8>], AccessError> {
1408                Err(AccessError::Read("test write failure".into()))
1409            }
1410        }
1411
1412        // SAFETY: The controller never exposes its inaccessible storage as initialized memory.
1413        let bytes =
1414            unsafe { Bytes::from_controller(Box::new(FailingController), size_of::<f32>()) };
1415        let data = TensorData::from_bytes(bytes, [1], DType::F32);
1416
1417        assert!(matches!(
1418            data.try_view::<f32>(),
1419            Err(DataError::StorageAccess(AccessError::Read(reason))) if reason == "test read failure"
1420        ));
1421    }
1422
1423    #[test]
1424    fn try_mut_view_validates_storage() {
1425        let mut invalid_representation =
1426            TensorData::from_bytes_unchecked(Bytes::from_bytes_vec(vec![0; 3]), [1], DType::F32);
1427        assert!(matches!(
1428            invalid_representation.try_mut_view::<f32>(),
1429            Err(DataError::InvalidRepresentation(_))
1430        ));
1431    }
1432
1433    #[test]
1434    fn try_cast_propagates_invalid_storage() {
1435        let invalid_inplace =
1436            TensorData::from_bytes_unchecked(Bytes::from_bytes_vec(vec![0; 3]), [1], DType::F32);
1437        assert!(matches!(
1438            invalid_inplace.try_cast(DType::I32),
1439            Err(DataError::InvalidRepresentation(_))
1440        ));
1441
1442        let invalid_clone =
1443            TensorData::from_bytes_unchecked(Bytes::from_bytes_vec(vec![0; 3]), [1], DType::F32);
1444        assert!(matches!(
1445            invalid_clone.try_cast(DType::F64),
1446            Err(DataError::InvalidRepresentation(_))
1447        ));
1448
1449        let invalid_count =
1450            TensorData::from_bytes_unchecked(Bytes::from_bytes_vec(vec![0; 8]), [1], DType::F32);
1451        assert!(matches!(
1452            invalid_count.try_cast(DType::F64),
1453            Err(DataError::ElementCountMismatch {
1454                expected: 1,
1455                actual: 2
1456            })
1457        ));
1458    }
1459
1460    #[test]
1461    fn try_cast_reports_unsupported_quantized_conversion() {
1462        let scheme = QuantScheme::default();
1463        let target = DType::QFloat(scheme);
1464
1465        assert_eq!(
1466            TensorData::from([1.0f32]).try_cast(target),
1467            Err(DataError::UnsupportedConversion {
1468                from: DType::F32,
1469                to: target,
1470            })
1471        );
1472
1473        let quantized = TensorData::quantized(vec![0i8], [1], scheme, &[1.0], None);
1474        assert_eq!(
1475            quantized.try_cast(DType::F32),
1476            Err(DataError::UnsupportedConversion {
1477                from: target,
1478                to: DType::F32,
1479            })
1480        );
1481    }
1482
1483    #[test]
1484    #[should_panic(expected = "Expected data type")]
1485    fn test_tensor_data_expect_mut_view_dtype_mismatch() {
1486        let mut data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]);
1487        let _view = data.mut_view::<i32>();
1488    }
1489
1490    #[test]
1491    fn test_tensor_data_index_view() {
1492        let data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]);
1493        let view = data.view::<f64>();
1494
1495        assert_eq!(view.shape(), &data.shape);
1496
1497        assert_eq!(view[&[0, 0]], 1.0);
1498        assert_eq!(view[&[0, 1]], 2.0);
1499        assert_eq!(view[&[1, 0]], 3.0);
1500        assert_eq!(view[&[1, 1]], 4.0);
1501    }
1502
1503    #[test]
1504    fn test_tensor_data_index_mut_view() {
1505        let mut data = TensorData::from([[1.0, 2.0], [3.0, 4.0]]);
1506        let shape = data.shape.clone();
1507
1508        let mut view = data.mut_view::<f64>();
1509
1510        assert_eq!(view.shape(), &shape);
1511
1512        assert_eq!(view[&[0, 0]], 1.0);
1513        assert_eq!(view[&[0, 1]], 2.0);
1514        assert_eq!(view[&[1, 0]], 3.0);
1515        assert_eq!(view[&[1, 1]], 4.0);
1516
1517        view[&[0, 0]] = 10.0;
1518        assert_eq!(view[&[0, 0]], 10.0);
1519    }
1520
1521    #[test]
1522    fn test_to_vec_as() {
1523        let data = TensorData::from([0.0f32, 1.0, 2.5]);
1524
1525        // Same-dtype copy.
1526        assert_eq!(data.try_to_vec_as::<f32>().unwrap(), vec![0.0f32, 1.0, 2.5]);
1527
1528        // Widening cast (different element size).
1529        assert_eq!(data.try_to_vec_as::<f64>().unwrap(), vec![0.0f64, 1.0, 2.5]);
1530
1531        // Float to int cast (same element size) truncates.
1532        assert_eq!(data.try_to_vec_as::<i32>().unwrap(), vec![0i32, 1, 2]);
1533
1534        // The source data is borrowed, not consumed.
1535        data.assert_eq(&TensorData::from([0.0f32, 1.0, 2.5]), true);
1536    }
1537
1538    #[test]
1539    fn test_into_vec_as() {
1540        let data = TensorData::from([0i32, 1, 2, 3]);
1541
1542        // Same-dtype conversion.
1543        assert_eq!(
1544            data.clone().try_into_vec_as::<i32>().unwrap(),
1545            vec![0i32, 1, 2, 3]
1546        );
1547
1548        // Int to float cast.
1549        assert_eq!(
1550            data.clone().try_into_vec_as::<f32>().unwrap(),
1551            vec![0.0f32, 1.0, 2.0, 3.0]
1552        );
1553
1554        // Narrowing int cast.
1555        assert_eq!(data.try_into_vec_as::<u8>().unwrap(), vec![0u8, 1, 2, 3]);
1556    }
1557}
burn_std/data/tensor/base.rs

burn_std/data/tensor/
base.rs
