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

burn_std/data/tensor/
base.rs
