burn_std/data/
tensor.rs

1use core::f32;
2
3use alloc::boxed::Box;
4use alloc::format;
5use alloc::string::String;
6use alloc::vec::Vec;
7use bytemuck::{AnyBitPattern, CheckedBitPattern, Zeroable, cast_mut, checked::CheckedCastError};
8use rand::Rng;
9use thiserror::Error;
10
11use crate::Scalar;
12use crate::distribution::Distribution;
13use crate::element::{Element, ElementConversion};
14use crate::tensor::DType;
15use crate::{
16    BoolStore, Bytes, QuantLevel, QuantMode, QuantScheme, QuantValue, QuantizedBytes, Shape, bf16,
17    f16,
18};
19
20use serde::{Deserialize, Serialize};
21
22/// Data structure for tensors.
23#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
24pub struct TensorData {
25    /// The values of the tensor (as bytes).
26    pub bytes: Bytes,
27
28    /// The shape of the tensor.
29    #[serde(with = "shape_inner")]
30    pub shape: Shape,
31
32    /// The data type of the tensor.
33    pub dtype: DType,
34}
35
36// For backward compatibility with shape `Vec<usize>`
37mod shape_inner {
38    use crate::SmallVec;
39
40    use super::*;
41
42    pub fn serialize<S: serde::Serializer>(
43        shape: &Shape,
44        serializer: S,
45    ) -> Result<S::Ok, S::Error> {
46        shape.as_slice().serialize(serializer)
47    }
48
49    pub fn deserialize<'de, D: serde::Deserializer<'de>>(
50        deserializer: D,
51    ) -> Result<Shape, D::Error> {
52        let dims = SmallVec::<[usize; _]>::deserialize(deserializer)?;
53        Ok(Shape::new_raw(dims))
54    }
55}
56
57impl TensorData {
58    /// Creates a new tensor data structure.
59    pub fn new<E: Element, S: Into<Shape>>(value: Vec<E>, shape: S) -> Self {
60        // Ensure shape is valid
61        let shape = shape.into();
62        Self::check_data_len(&value, &shape);
63
64        Self {
65            bytes: Bytes::from_elems(value),
66            shape,
67            dtype: E::dtype(),
68        }
69    }
70
71    /// Creates a new quantized tensor data structure.
72    pub fn quantized<E: Element, S: Into<Shape>>(
73        value: Vec<E>,
74        shape: S,
75        scheme: QuantScheme,
76        qparams: &[f32],
77    ) -> Self {
78        let shape = shape.into();
79        Self::check_data_len(&value, &shape);
80
81        let q_bytes = QuantizedBytes::new(value, scheme, qparams);
82
83        Self {
84            bytes: q_bytes.bytes,
85            shape,
86            dtype: DType::QFloat(q_bytes.scheme),
87        }
88    }
89
90    /// Creates a new tensor data structure from raw bytes.
91    pub fn from_bytes<S: Into<Shape>>(bytes: Bytes, shape: S, dtype: DType) -> Self {
92        Self {
93            bytes,
94            shape: shape.into(),
95            dtype,
96        }
97    }
98
99    /// Creates a new tensor data structure from raw bytes stored in a vector.
100    ///
101    /// Prefer [`TensorData::new`] or [`TensorData::quantized`] over this method unless you are
102    /// certain that the bytes representation is valid.
103    pub fn from_bytes_vec<S: Into<Shape>>(bytes: Vec<u8>, shape: S, dtype: DType) -> Self {
104        Self {
105            bytes: Bytes::from_bytes_vec(bytes),
106            shape: shape.into(),
107            dtype,
108        }
109    }
110
111    // Check that the input vector contains a correct number of elements
112    fn check_data_len<E: Element>(data: &[E], shape: &Shape) {
113        let expected_data_len = Self::numel(shape);
114        let num_data = data.len();
115        assert_eq!(
116            expected_data_len, num_data,
117            "Shape {shape:?} is invalid for input of size {num_data:?}",
118        );
119    }
120
121    /// Returns the immutable slice view of the tensor data.
122    pub fn as_slice<E: Element>(&self) -> Result<&[E], DataError> {
123        if self.matches_target_dtype::<E>() {
124            match E::dtype() {
125                // The only way to create a bool `TensorData` with invalid values is by unsafely modifying
126                // the dtype. This should be considered unsafe to begin with, so we unsafely cast bool
127                // to u8 to skip bit validation. Validation iterates through the entire vector, so it's slow.
128                DType::Bool(BoolStore::Native) => {
129                    let slice = bytemuck::checked::try_cast_slice::<_, u8>(&self.bytes)
130                        .map_err(DataError::CastError)?;
131                    Ok(unsafe { core::mem::transmute::<&[u8], &[E]>(slice) })
132                }
133                _ => bytemuck::checked::try_cast_slice(&self.bytes).map_err(DataError::CastError),
134            }
135        } else {
136            Err(DataError::TypeMismatch(format!(
137                "Invalid target element type (expected {:?}, got {:?})",
138                self.dtype,
139                E::dtype()
140            )))
141        }
142    }
143
144    /// Returns the mutable slice view of the tensor data.
145    ///
146    /// # Panics
147    /// If the target element type is different from the stored element type.
148    pub fn as_mut_slice<E: Element>(&mut self) -> Result<&mut [E], DataError> {
149        if self.matches_target_dtype::<E>() {
150            match E::dtype() {
151                // The only way to create a bool `TensorData` with invalid values is by unsafely modifying
152                // the dtype. This should be considered unsafe to begin with, so we unsafely cast bool
153                // to u8 to skip bit validation. Validation iterates through the entire vector, so it's slow.
154                DType::Bool(BoolStore::Native) => {
155                    let slice = bytemuck::checked::try_cast_slice_mut::<_, u8>(&mut self.bytes)
156                        .map_err(DataError::CastError)?;
157                    Ok(unsafe { core::mem::transmute::<&mut [u8], &mut [E]>(slice) })
158                }
159                _ => bytemuck::checked::try_cast_slice_mut(&mut self.bytes)
160                    .map_err(DataError::CastError),
161            }
162        } else {
163            Err(DataError::TypeMismatch(format!(
164                "Invalid target element type (expected {:?}, got {:?})",
165                self.dtype,
166                E::dtype()
167            )))
168        }
169    }
170
171    /// Returns the tensor data as a vector of scalar values.
172    pub fn to_vec<E: Element>(&self) -> Result<Vec<E>, DataError> {
173        Ok(self.as_slice()?.to_vec())
174    }
175
176    /// Returns the tensor data as a vector of scalar values.
177    pub fn into_vec<E: Element>(self) -> Result<Vec<E>, DataError> {
178        // This means we cannot call `into_vec` for QFloat
179        if !self.matches_target_dtype::<E>() {
180            return Err(DataError::TypeMismatch(format!(
181                "Invalid target element type (expected {:?}, got {:?})",
182                self.dtype,
183                E::dtype()
184            )));
185        }
186
187        match E::dtype() {
188            // The only way to create a bool `TensorData` with invalid values is by unsafely modifying
189            // the dtype. This should be considered unsafe to begin with, so we unsafely cast bool
190            // to u8 to skip bit validation. Validation iterates through the entire vector, so it's slow.
191            DType::Bool(BoolStore::Native) => {
192                let vec = self.into_vec_unchecked::<u8>()?;
193                Ok(unsafe { core::mem::transmute::<Vec<u8>, Vec<E>>(vec) })
194            }
195            _ => self.into_vec_unchecked(),
196        }
197    }
198
199    /// Returns the tensor data as a vector of scalar values. Does not check dtype.
200    fn into_vec_unchecked<E: Element>(self) -> Result<Vec<E>, DataError> {
201        let mut me = self;
202        me.bytes = match me.bytes.try_into_vec::<E>() {
203            Ok(elems) => return Ok(elems),
204            Err(bytes) => bytes,
205        };
206
207        // The bytes might have been deserialized and allocated with a different align.
208        // In that case, we have to memcopy the data into a new vector, more suitably allocated
209        Ok(bytemuck::checked::try_cast_slice(me.as_bytes())
210            .map_err(DataError::CastError)?
211            .to_vec())
212    }
213
214    fn matches_target_dtype<E: Element>(&self) -> bool {
215        let target_dtype = E::dtype();
216        match self.dtype {
217            DType::Bool(BoolStore::U8) => {
218                matches!(target_dtype, DType::U8 | DType::Bool(BoolStore::U8))
219            }
220            DType::Bool(BoolStore::U32) => {
221                matches!(target_dtype, DType::U32 | DType::Bool(BoolStore::U32))
222            }
223            dtype => dtype == target_dtype,
224        }
225    }
226
227    /// Returns an iterator over the values of the tensor data.
228    pub fn iter<E: Element>(&self) -> Box<dyn Iterator<Item = E> + '_> {
229        if E::dtype() == self.dtype {
230            Box::new(bytemuck::checked::cast_slice(&self.bytes).iter().copied())
231        } else {
232            match self.dtype {
233                DType::I8 => Box::new(
234                    bytemuck::checked::cast_slice(&self.bytes)
235                        .iter()
236                        .map(|e: &i8| e.elem::<E>()),
237                ),
238                DType::I16 => Box::new(
239                    bytemuck::checked::cast_slice(&self.bytes)
240                        .iter()
241                        .map(|e: &i16| e.elem::<E>()),
242                ),
243                DType::I32 => Box::new(
244                    bytemuck::checked::cast_slice(&self.bytes)
245                        .iter()
246                        .map(|e: &i32| e.elem::<E>()),
247                ),
248                DType::I64 => Box::new(
249                    bytemuck::checked::cast_slice(&self.bytes)
250                        .iter()
251                        .map(|e: &i64| e.elem::<E>()),
252                ),
253                DType::U8 => Box::new(self.bytes.iter().map(|e| e.elem::<E>())),
254                DType::U16 => Box::new(
255                    bytemuck::checked::cast_slice(&self.bytes)
256                        .iter()
257                        .map(|e: &u16| e.elem::<E>()),
258                ),
259                DType::U32 => Box::new(
260                    bytemuck::checked::cast_slice(&self.bytes)
261                        .iter()
262                        .map(|e: &u32| e.elem::<E>()),
263                ),
264                DType::U64 => Box::new(
265                    bytemuck::checked::cast_slice(&self.bytes)
266                        .iter()
267                        .map(|e: &u64| e.elem::<E>()),
268                ),
269                DType::BF16 => Box::new(
270                    bytemuck::checked::cast_slice(&self.bytes)
271                        .iter()
272                        .map(|e: &bf16| e.elem::<E>()),
273                ),
274                DType::F16 => Box::new(
275                    bytemuck::checked::cast_slice(&self.bytes)
276                        .iter()
277                        .map(|e: &f16| e.elem::<E>()),
278                ),
279                DType::F32 | DType::Flex32 => Box::new(
280                    bytemuck::checked::cast_slice(&self.bytes)
281                        .iter()
282                        .map(|e: &f32| e.elem::<E>()),
283                ),
284                DType::F64 => Box::new(
285                    bytemuck::checked::cast_slice(&self.bytes)
286                        .iter()
287                        .map(|e: &f64| e.elem::<E>()),
288                ),
289                // bool is a byte value equal to either 0 or 1
290                DType::Bool(BoolStore::Native) | DType::Bool(BoolStore::U8) => {
291                    Box::new(self.bytes.iter().map(|e| e.elem::<E>()))
292                }
293                DType::Bool(BoolStore::U32) => Box::new(
294                    bytemuck::checked::cast_slice(&self.bytes)
295                        .iter()
296                        .map(|e: &u32| e.elem::<E>()),
297                ),
298                DType::QFloat(scheme) => match scheme {
299                    QuantScheme {
300                        level: QuantLevel::Tensor | QuantLevel::Block(_),
301                        mode: QuantMode::Symmetric,
302                        value:
303                            QuantValue::Q8F
304                            | QuantValue::Q8S
305                            // Represent sub-byte values as i8
306                            | QuantValue::Q4F
307                            | QuantValue::Q4S
308                            | QuantValue::Q2F
309                            | QuantValue::Q2S,
310                        ..
311                    } => {
312                        // Quantized int8 values
313                        let q_bytes = QuantizedBytes {
314                            bytes: self.bytes.clone(),
315                            scheme,
316                            num_elements: self.num_elements(),
317                        };
318                        let (values, _) = q_bytes.into_vec_i8();
319
320                        Box::new(
321                            values
322                                .iter()
323                                .map(|e: &i8| e.elem::<E>())
324                                .collect::<Vec<_>>()
325                                .into_iter(),
326                        )
327                    }
328                    QuantScheme {
329                        level: QuantLevel::Tensor | QuantLevel::Block(_),
330                        mode: QuantMode::Symmetric,
331                        value:
332                            QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1,
333                        ..
334                    } => {
335                        unimplemented!("Not yet implemented for iteration");
336                    }
337                    QuantScheme {
338                        level: QuantLevel::BlockTensor { .. },
339                        ..
340                    } => {
341                        unimplemented!("two-level quantization is not supported yet")
342                    }
343                },
344            }
345        }
346    }
347
348    /// Returns the rank (the number of dimensions).
349    pub fn rank(&self) -> usize {
350        self.shape.len()
351    }
352
353    /// Returns the total number of elements of the tensor data.
354    pub fn num_elements(&self) -> usize {
355        Self::numel(&self.shape)
356    }
357
358    fn numel(shape: &[usize]) -> usize {
359        shape.iter().product()
360    }
361
362    /// Populates the data with random values.
363    pub fn random<E: Element, R: Rng, S: Into<Shape>>(
364        shape: S,
365        distribution: Distribution,
366        rng: &mut R,
367    ) -> Self {
368        let shape = shape.into();
369        let num_elements = Self::numel(&shape);
370        let mut data = Vec::with_capacity(num_elements);
371
372        for _ in 0..num_elements {
373            data.push(E::random(distribution, rng));
374        }
375
376        TensorData::new(data, shape)
377    }
378
379    /// Populates the data with zeros.
380    pub fn zeros<E: Element, S: Into<Shape>>(shape: S) -> TensorData {
381        let shape = shape.into();
382        let num_elements = Self::numel(&shape);
383        let mut data = Vec::<E>::with_capacity(num_elements);
384
385        for _ in 0..num_elements {
386            data.push(0.elem());
387        }
388
389        TensorData::new(data, shape)
390    }
391
392    /// Populates the data with ones.
393    pub fn ones<E: Element, S: Into<Shape>>(shape: S) -> TensorData {
394        let shape = shape.into();
395        let num_elements = Self::numel(&shape);
396        let mut data = Vec::<E>::with_capacity(num_elements);
397
398        for _ in 0..num_elements {
399            data.push(1.elem());
400        }
401
402        TensorData::new(data, shape)
403    }
404
405    /// Populates the data with the given value
406    pub fn full<E: Element, S: Into<Shape>>(shape: S, fill_value: E) -> TensorData {
407        let shape = shape.into();
408        let num_elements = Self::numel(&shape);
409        let mut data = Vec::<E>::with_capacity(num_elements);
410        for _ in 0..num_elements {
411            data.push(fill_value)
412        }
413
414        TensorData::new(data, shape)
415    }
416
417    /// Populates the data with the given value
418    pub fn full_dtype<E: Into<Scalar>, S: Into<Shape>>(
419        shape: S,
420        fill_value: E,
421        dtype: DType,
422    ) -> TensorData {
423        let fill_value = fill_value.into();
424        match dtype {
425            DType::F64 => Self::full::<f64, _>(shape, fill_value.elem()),
426            DType::F32 | DType::Flex32 => Self::full::<f32, _>(shape, fill_value.elem()),
427            DType::F16 => Self::full::<f16, _>(shape, fill_value.elem()),
428            DType::BF16 => Self::full::<bf16, _>(shape, fill_value.elem()),
429            DType::I64 => Self::full::<i64, _>(shape, fill_value.elem()),
430            DType::I32 => Self::full::<i32, _>(shape, fill_value.elem()),
431            DType::I16 => Self::full::<i16, _>(shape, fill_value.elem()),
432            DType::I8 => Self::full::<i8, _>(shape, fill_value.elem()),
433            DType::U64 => Self::full::<u64, _>(shape, fill_value.elem()),
434            DType::U32 => Self::full::<u32, _>(shape, fill_value.elem()),
435            DType::U16 => Self::full::<u16, _>(shape, fill_value.elem()),
436            DType::U8 => Self::full::<u8, _>(shape, fill_value.elem()),
437            DType::Bool(BoolStore::Native) => Self::full::<bool, _>(shape, fill_value.elem()),
438            DType::Bool(BoolStore::U8) => {
439                Self::full::<u8, _>(shape, fill_value.elem()).into_bool_u8()
440            }
441            DType::Bool(BoolStore::U32) => {
442                Self::full::<u32, _>(shape, fill_value.elem()).into_bool_u32()
443            }
444            DType::QFloat(_) => unreachable!(),
445        }
446    }
447
448    // Unchecked, used to overwrite the dtype
449    fn into_bool_u8(mut self) -> Self {
450        self.dtype = DType::Bool(BoolStore::U8);
451        self
452    }
453
454    // Unchecked, used to overwrite the dtype
455    fn into_bool_u32(mut self) -> Self {
456        self.dtype = DType::Bool(BoolStore::U32);
457        self
458    }
459
460    /// Converts the data to a different element type.
461    pub fn convert<E: Element>(self) -> Self {
462        self.convert_dtype(E::dtype())
463    }
464
465    /// Converts the data to a different element type.
466    pub fn convert_dtype(self, dtype: DType) -> Self {
467        if dtype == self.dtype {
468            self
469        } else if dtype.size() == self.dtype.size()
470            && !matches!(
471                self.dtype,
472                DType::Bool(BoolStore::Native) | DType::QFloat(_)
473            )
474            && !matches!(dtype, DType::Bool(BoolStore::Native) | DType::QFloat(_))
475        {
476            match self.dtype {
477                DType::F64 => self.convert_inplace_dtype::<f64>(dtype),
478                DType::F32 | DType::Flex32 => self.convert_inplace_dtype::<f32>(dtype),
479                DType::F16 => self.convert_inplace_dtype::<f16>(dtype),
480                DType::BF16 => self.convert_inplace_dtype::<bf16>(dtype),
481                DType::I64 => self.convert_inplace_dtype::<i64>(dtype),
482                DType::I32 => self.convert_inplace_dtype::<i32>(dtype),
483                DType::I16 => self.convert_inplace_dtype::<i16>(dtype),
484                DType::I8 => self.convert_inplace_dtype::<i8>(dtype),
485                DType::U64 => self.convert_inplace_dtype::<u64>(dtype),
486                DType::U32 => self.convert_inplace_dtype::<u32>(dtype),
487                DType::U16 => self.convert_inplace_dtype::<u16>(dtype),
488                DType::U8 => self.convert_inplace_dtype::<u8>(dtype),
489                DType::Bool(BoolStore::U8) => self.convert_inplace_dtype::<u8>(dtype),
490                DType::Bool(BoolStore::U32) => self.convert_inplace_dtype::<u32>(dtype),
491                DType::Bool(BoolStore::Native) | DType::QFloat(_) => unreachable!(),
492            }
493        } else {
494            match self.dtype {
495                DType::F64 => self.convert_clone_dtype::<f64>(dtype),
496                DType::F32 | DType::Flex32 => self.convert_clone_dtype::<f32>(dtype),
497                DType::F16 => self.convert_clone_dtype::<f16>(dtype),
498                DType::BF16 => self.convert_clone_dtype::<bf16>(dtype),
499                DType::I64 => self.convert_clone_dtype::<i64>(dtype),
500                DType::I32 => self.convert_clone_dtype::<i32>(dtype),
501                DType::I16 => self.convert_clone_dtype::<i16>(dtype),
502                DType::I8 => self.convert_clone_dtype::<i8>(dtype),
503                DType::U64 => self.convert_clone_dtype::<u64>(dtype),
504                DType::U32 => self.convert_clone_dtype::<u32>(dtype),
505                DType::U16 => self.convert_clone_dtype::<u16>(dtype),
506                DType::U8 => self.convert_clone_dtype::<u8>(dtype),
507                DType::Bool(BoolStore::Native) => self.convert_clone_dtype::<bool>(dtype),
508                DType::Bool(BoolStore::U8) => self.convert_clone_dtype::<u8>(dtype),
509                DType::Bool(BoolStore::U32) => self.convert_clone_dtype::<u32>(dtype),
510                DType::QFloat(_) => unreachable!(),
511            }
512        }
513    }
514
515    fn convert_inplace_dtype<Current: Element + AnyBitPattern>(self, dtype: DType) -> Self {
516        match dtype {
517            DType::F64 => self.convert_inplace::<Current, f64>(),
518            DType::F32 | DType::Flex32 => self.convert_inplace::<Current, f32>(),
519            DType::F16 => self.convert_inplace::<Current, f16>(),
520            DType::BF16 => self.convert_inplace::<Current, bf16>(),
521            DType::I64 => self.convert_inplace::<Current, i64>(),
522            DType::I32 => self.convert_inplace::<Current, i32>(),
523            DType::I16 => self.convert_inplace::<Current, i16>(),
524            DType::I8 => self.convert_inplace::<Current, i8>(),
525            DType::U64 => self.convert_inplace::<Current, u64>(),
526            DType::U32 => self.convert_inplace::<Current, u32>(),
527            DType::U16 => self.convert_inplace::<Current, u16>(),
528            DType::U8 => self.convert_inplace::<Current, u8>(),
529            DType::Bool(BoolStore::U8) => self.convert_inplace_bool::<Current, u8>().into_bool_u8(),
530            DType::Bool(BoolStore::U32) => {
531                self.convert_inplace_bool::<Current, u32>().into_bool_u32()
532            }
533            DType::Bool(BoolStore::Native) | DType::QFloat(_) => unreachable!(),
534        }
535    }
536
537    fn convert_inplace<Current: Element + AnyBitPattern, Target: Element + AnyBitPattern>(
538        self,
539    ) -> Self {
540        self.convert_inplace_with::<Current, Target>(|x| x.elem())
541    }
542
543    fn convert_inplace_bool<Current: Element + AnyBitPattern, Target: Element + AnyBitPattern>(
544        self,
545    ) -> Self {
546        self.convert_inplace_with::<Current, Target>(|x| x.to_bool().elem())
547    }
548
549    fn convert_inplace_with<Current: Element + AnyBitPattern, Target: Element + AnyBitPattern>(
550        mut self,
551        transform: impl Fn(&Current) -> Target,
552    ) -> Self {
553        for x in bytemuck::cast_slice_mut::<_, Current>(&mut self.bytes) {
554            let t = transform(x);
555            let x = cast_mut::<_, Target>(x);
556            *x = t;
557        }
558
559        self.dtype = Target::dtype();
560
561        self
562    }
563
564    fn convert_clone_dtype<Current: Element + CheckedBitPattern>(self, dtype: DType) -> Self {
565        match dtype {
566            DType::F64 => self.convert_clone::<Current, f64>(),
567            DType::F32 | DType::Flex32 => self.convert_clone::<Current, f32>(),
568            DType::F16 => self.convert_clone::<Current, f16>(),
569            DType::BF16 => self.convert_clone::<Current, bf16>(),
570            DType::I64 => self.convert_clone::<Current, i64>(),
571            DType::I32 => self.convert_clone::<Current, i32>(),
572            DType::I16 => self.convert_clone::<Current, i16>(),
573            DType::I8 => self.convert_clone::<Current, i8>(),
574            DType::U64 => self.convert_clone::<Current, u64>(),
575            DType::U32 => self.convert_clone::<Current, u32>(),
576            DType::U16 => self.convert_clone::<Current, u16>(),
577            DType::U8 => self.convert_clone::<Current, u8>(),
578            DType::Bool(BoolStore::Native) => self.convert_clone::<Current, bool>(),
579            DType::Bool(BoolStore::U8) => self.convert_clone_bool::<Current, u8>().into_bool_u8(),
580            DType::Bool(BoolStore::U32) => {
581                self.convert_clone_bool::<Current, u32>().into_bool_u32()
582            }
583            DType::QFloat(_) => unreachable!(),
584        }
585    }
586
587    fn convert_clone<Current: Element + CheckedBitPattern, Target: Element + Zeroable>(
588        self,
589    ) -> Self {
590        self.convert_clone_with::<Current, Target>(|x| x.elem())
591    }
592
593    fn convert_clone_bool<Current: Element + CheckedBitPattern, Target: Element + Zeroable>(
594        self,
595    ) -> Self {
596        self.convert_clone_with::<Current, Target>(|x| x.to_bool().elem())
597    }
598
599    fn convert_clone_with<Current: Element + CheckedBitPattern, Target: Element + Zeroable>(
600        self,
601        transform: impl Fn(&Current) -> Target,
602    ) -> Self {
603        let this = bytemuck::checked::cast_slice::<_, Current>(&self.bytes);
604        let mut out: Vec<Target> = ::alloc::vec![Zeroable::zeroed(); self.num_elements()];
605
606        for (x, out) in this.iter().zip(&mut out) {
607            *out = transform(x);
608        }
609
610        Self::new(out, self.shape)
611    }
612
613    /// Returns the data as a slice of bytes.
614    pub fn as_bytes(&self) -> &[u8] {
615        &self.bytes
616    }
617
618    /// Returns the bytes representation of the data.
619    pub fn into_bytes(self) -> Bytes {
620        self.bytes
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                    level: QuantLevel::Tensor | QuantLevel::Block(_),
751                    mode: QuantMode::Symmetric,
752                    value:
753                        QuantValue::Q8F
754                        | QuantValue::Q8S
755                        // Display sub-byte values as i8
756                        | QuantValue::Q4F
757                        | QuantValue::Q4S
758                        | QuantValue::Q2F
759                        | QuantValue::Q2S,
760                    ..
761                } => {
762                    format!("{:?} {scheme:?}", self.iter::<i8>().collect::<Vec<_>>())
763                },
764                QuantScheme {
765                        level: QuantLevel::Tensor | QuantLevel::Block(_),
766                        mode: QuantMode::Symmetric,
767                        value:
768                            QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1,
769                        ..
770                    } => {
771                        unimplemented!("Can't format yet");
772                    }
773                QuantScheme {
774                    level: QuantLevel::BlockTensor { .. },
775                    ..
776                } => {
777                    unimplemented!("two-level quantization is not supported yet")
778                }
779            },
780        };
781        f.write_str(fmt.as_str())
782    }
783}
784
785/// The things that can go wrong when manipulating tensor data.
786#[derive(Debug, Error)]
787pub enum DataError {
788    /// Failed to cast the values to a specified element type.
789    #[error("Failed to cast values to the specified element type.\nError:\n  {0}")]
790    CastError(CheckedCastError),
791    /// Invalid target element type.
792    #[error("{0}")]
793    TypeMismatch(String),
794}
795
796#[cfg(test)]
797mod tests {
798    use super::*;
799    use crate::shape;
800    use alloc::vec;
801    use rand::{
802        SeedableRng,
803        rngs::{StdRng, SysRng},
804    };
805
806    #[test]
807    fn should_have_rank() {
808        let shape = [3, 5, 6];
809        let data = TensorData::random::<f32, _, _>(
810            shape,
811            Distribution::Default,
812            &mut StdRng::try_from_rng(&mut SysRng).unwrap(),
813        );
814
815        assert_eq!(data.rank(), 3);
816    }
817
818    #[test]
819    fn into_vec_should_yield_same_value_as_iter() {
820        let shape = [3, 5, 6];
821        let data = TensorData::random::<f32, _, _>(
822            shape,
823            Distribution::Default,
824            &mut StdRng::try_from_rng(&mut SysRng).unwrap(),
825        );
826
827        let expected = data.iter::<f32>().collect::<Vec<f32>>();
828        let actual = data.into_vec::<f32>().unwrap();
829
830        assert_eq!(expected, actual);
831    }
832
833    #[test]
834    #[should_panic]
835    fn into_vec_should_assert_wrong_dtype() {
836        let shape = [3, 5, 6];
837        let data = TensorData::random::<f32, _, _>(
838            shape,
839            Distribution::Default,
840            &mut StdRng::try_from_rng(&mut SysRng).unwrap(),
841        );
842
843        data.into_vec::<i32>().unwrap();
844    }
845
846    #[test]
847    fn should_have_right_num_elements() {
848        let shape = [3, 5, 6];
849        let num_elements: usize = shape.iter().product();
850        let data = TensorData::random::<f32, _, _>(
851            shape,
852            Distribution::Default,
853            &mut StdRng::try_from_rng(&mut SysRng).unwrap(),
854        );
855
856        assert_eq!(num_elements, data.bytes.len() / 4); // f32 stored as u8s
857        assert_eq!(num_elements, data.as_slice::<f32>().unwrap().len());
858    }
859
860    #[test]
861    fn should_have_right_shape() {
862        let data = TensorData::from([[3.0, 5.0, 6.0]]);
863        assert_eq!(data.shape, shape![1, 3]);
864
865        let data = TensorData::from([[4.0, 5.0, 8.0], [3.0, 5.0, 6.0]]);
866        assert_eq!(data.shape, shape![2, 3]);
867
868        let data = TensorData::from([3.0, 5.0, 6.0]);
869        assert_eq!(data.shape, shape![3]);
870    }
871
872    #[test]
873    fn should_convert_bytes_correctly() {
874        let mut vector: Vec<f32> = Vec::with_capacity(5);
875        vector.push(2.0);
876        vector.push(3.0);
877        let data1 = TensorData::new(vector, vec![2]);
878
879        let factor = core::mem::size_of::<f32>() / core::mem::size_of::<u8>();
880        assert_eq!(data1.bytes.len(), 2 * factor);
881        assert_eq!(data1.bytes.capacity(), 5 * factor);
882    }
883
884    #[test]
885    fn should_convert_bytes_correctly_inplace() {
886        fn test_precision<E: Element>() {
887            let data = TensorData::new((0..32).collect(), [32]);
888            for (i, val) in data
889                .clone()
890                .convert::<E>()
891                .into_vec::<E>()
892                .unwrap()
893                .into_iter()
894                .enumerate()
895            {
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_convert_negative_values_to_bool_store() {
907        for store in [BoolStore::U8, BoolStore::U32, BoolStore::Native] {
908            let data = TensorData::from([-1i32, 0, 1, -12]).convert_dtype(DType::Bool(store));
909            assert_eq!(data.dtype, DType::Bool(store));
910            assert_eq!(
911                data.iter::<bool>().collect::<Vec<_>>(),
912                [true, false, true, true]
913            );
914
915            let data = TensorData::from([-1.5f32, 0.0, 0.5]).convert_dtype(DType::Bool(store));
916            assert_eq!(data.iter::<bool>().collect::<Vec<_>>(), [true, false, true]);
917        }
918    }
919
920    macro_rules! test_dtypes {
921    ($test_name:ident, $($dtype:ty),*) => {
922        $(
923            paste::paste! {
924                #[test]
925                fn [<$test_name _ $dtype:snake>]() {
926                    let full_dtype = TensorData::full_dtype([2, 16], 4, <$dtype>::dtype());
927                    let full = TensorData::full::<$dtype, _>([2, 16], 4.elem());
928                    assert_eq!(full_dtype, full);
929                }
930            }
931        )*
932    };
933}
934
935    test_dtypes!(
936        should_create_with_dtype,
937        bool,
938        i8,
939        i16,
940        i32,
941        i64,
942        u8,
943        u16,
944        u32,
945        u64,
946        f16,
947        bf16,
948        f32,
949        f64
950    );
951
952    #[test]
953    fn should_serialize_deserialize_tensor_data() {
954        let data = TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3]);
955        assert_eq!(
956            data.as_bytes(),
957            [
958                0, 0, 128, 63, 0, 0, 0, 64, 0, 0, 64, 64, 0, 0, 128, 64, 0, 0, 160, 64, 0, 0, 192,
959                64
960            ]
961        );
962        let serialized = serde_json::to_string(&data).unwrap();
963        let deserialized: TensorData = serde_json::from_str(&serialized).unwrap();
964        assert_eq!(data, deserialized);
965    }
966
967    #[test]
968    fn should_deserialize_tensor_data_with_shape_inner() {
969        // TensorData `shape` was previously a Vec<usize>.
970        let serialized = r#"{
971        "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],
972        "shape": [2, 3],
973        "dtype": "F32"
974    }"#;
975
976        let data: TensorData = serde_json::from_str(serialized).unwrap();
977        assert_eq!(data.shape, shape![2, 3]);
978        assert_eq!(
979            data.as_slice::<f32>().unwrap(),
980            &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
981        );
982    }
983
984    #[test]
985    fn should_serialize_shape_as_flat_array() {
986        // Ensure the new Shape serializes identically to how Vec<usize> used to,
987        // i.e. as a flat JSON array, not as an object like `{"dims": [2, 3]}`.
988        let data = TensorData::new(vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3]);
989        let serialized = serde_json::to_string(&data).unwrap();
990        let json: serde_json::Value = serde_json::from_str(&serialized).unwrap();
991        assert_eq!(json["shape"], serde_json::json!([2, 3]));
992    }
993}
burn_std/data/tensor.rs

burn_std/data/
tensor.rs
