burn_std/data/tensor/
conversion.rs

1use alloc::boxed::Box;
2use alloc::vec;
3use alloc::vec::Vec;
4
5use bytemuck::{AnyBitPattern, CheckedBitPattern, Zeroable, cast_mut};
6
7use crate::element::{Element, ElementConversion};
8use crate::tensor::DType;
9use crate::{
10    BoolStore, QuantMode, QuantScheme, QuantValue, QuantizedBytes, Reader, Writer, bf16, f16,
11};
12
13use super::{DataError, TensorData};
14
15impl TensorData {
16    /// Copies and converts the data to a [`Vec<E>`].
17    ///
18    /// By contract, this is equivalent to:
19    /// `data.clone().try_into_vec_as::<E>()`
20    ///
21    /// Particular conversions may provide more efficient implementations.
22    ///
23    /// # Errors
24    ///
25    /// Returns an error if storage access fails, the conversion isn't supported, or the stored
26    /// representation or element count is invalid.
27    pub fn try_to_vec_as<E: Element>(&self) -> Result<Vec<E>, DataError> {
28        self.clone().try_into_vec_as::<E>()
29    }
30
31    /// Converts the data to a [`Vec<E>`].
32    ///
33    /// By contract, this is equivalent to:
34    /// `data.try_cast_as::<E>()?.try_into_vec::<E>()`
35    ///
36    /// Particular conversions may provide more efficient implementations.
37    ///
38    /// # Errors
39    ///
40    /// Returns an error if storage access fails, the conversion isn't supported, or the stored
41    /// representation or element count is invalid.
42    pub fn try_into_vec_as<E: Element>(self) -> Result<Vec<E>, DataError> {
43        self.try_cast_as::<E>()?.into_vec_unchecked::<E>()
44    }
45
46    /// Copies the stored values to a vector without dtype conversion.
47    #[deprecated(since = "0.22.0", note = "use try_to_vec::<E>()")]
48    pub fn to_vec<E: Element>(&self) -> Result<Vec<E>, DataError> {
49        self.try_to_vec::<E>()
50    }
51
52    /// Converts the stored values into a vector without dtype conversion.
53    #[deprecated(since = "0.22.0", note = "use try_into_vec::<E>()")]
54    pub fn into_vec<E: Element>(self) -> Result<Vec<E>, DataError> {
55        self.try_into_vec::<E>()
56    }
57
58    /// Copies the stored values to a vector without dtype conversion.
59    ///
60    /// # Errors
61    ///
62    /// Returns an error if storage access fails, the stored dtype doesn't match `E`, or the byte
63    /// representation is invalid for `E`.
64    pub fn try_to_vec<E: Element>(&self) -> Result<Vec<E>, DataError> {
65        Ok(self.as_slice()?.to_vec())
66    }
67
68    /// Converts the stored values into a vector without dtype conversion.
69    ///
70    /// This may reuse the underlying allocation when its layout and ownership permit it.
71    ///
72    /// # Errors
73    ///
74    /// Returns an error if storage access fails, the stored dtype doesn't match `E`, or the byte
75    /// representation is invalid for `E`.
76    pub fn try_into_vec<E: Element>(self) -> Result<Vec<E>, DataError> {
77        // This means we cannot call `into_vec` for QFloat
78        if !self.matches_target_dtype::<E>() {
79            return Err(DataError::dtype_mismatch_as::<E>(self.dtype));
80        }
81
82        self.into_vec_unchecked()
83    }
84
85    /// Returns the tensor data as a vector of scalar values. Does not check dtype.
86    fn into_vec_unchecked<E: Element>(self) -> Result<Vec<E>, DataError> {
87        let mut me = self;
88        me.bytes.read(Reader::new())?;
89        me.bytes = match me.bytes.try_into_vec::<E>() {
90            Ok(elems) => return Ok(elems),
91            Err(bytes) => bytes,
92        };
93
94        // The bytes might have been deserialized and allocated with a different align.
95        // In that case, we have to memcopy the data into a new vector, more suitably allocated
96        Ok(
97            bytemuck::checked::try_cast_slice(me.bytes.read(Reader::new())?)
98                .map_err(DataError::InvalidRepresentation)?
99                .to_vec(),
100        )
101    }
102
103    /// Returns an iterator over the values of the tensor data.
104    pub fn iter<E: Element>(&self) -> Box<dyn Iterator<Item = E> + '_> {
105        self.iter_exact::<E>()
106    }
107
108    /// Returns an exact-size iterator for internal data comparisons.
109    pub(crate) fn iter_exact<E: Element>(&self) -> Box<dyn ExactSizeIterator<Item = E> + '_> {
110        if E::dtype() == self.dtype {
111            Box::new(bytemuck::checked::cast_slice(&self.bytes).iter().copied())
112        } else {
113            match self.dtype {
114                DType::I8 => Box::new(
115                    bytemuck::checked::cast_slice(&self.bytes)
116                        .iter()
117                        .map(|e: &i8| e.elem::<E>()),
118                ),
119                DType::I16 => Box::new(
120                    bytemuck::checked::cast_slice(&self.bytes)
121                        .iter()
122                        .map(|e: &i16| e.elem::<E>()),
123                ),
124                DType::I32 => Box::new(
125                    bytemuck::checked::cast_slice(&self.bytes)
126                        .iter()
127                        .map(|e: &i32| e.elem::<E>()),
128                ),
129                DType::I64 => Box::new(
130                    bytemuck::checked::cast_slice(&self.bytes)
131                        .iter()
132                        .map(|e: &i64| e.elem::<E>()),
133                ),
134                DType::U8 => Box::new(self.bytes.iter().map(|e| e.elem::<E>())),
135                DType::U16 => Box::new(
136                    bytemuck::checked::cast_slice(&self.bytes)
137                        .iter()
138                        .map(|e: &u16| e.elem::<E>()),
139                ),
140                DType::U32 => Box::new(
141                    bytemuck::checked::cast_slice(&self.bytes)
142                        .iter()
143                        .map(|e: &u32| e.elem::<E>()),
144                ),
145                DType::U64 => Box::new(
146                    bytemuck::checked::cast_slice(&self.bytes)
147                        .iter()
148                        .map(|e: &u64| e.elem::<E>()),
149                ),
150                DType::BF16 => Box::new(
151                    bytemuck::checked::cast_slice(&self.bytes)
152                        .iter()
153                        .map(|e: &bf16| e.elem::<E>()),
154                ),
155                DType::F16 => Box::new(
156                    bytemuck::checked::cast_slice(&self.bytes)
157                        .iter()
158                        .map(|e: &f16| e.elem::<E>()),
159                ),
160                DType::F32 | DType::Flex32 => Box::new(
161                    bytemuck::checked::cast_slice(&self.bytes)
162                        .iter()
163                        .map(|e: &f32| e.elem::<E>()),
164                ),
165                DType::F64 => Box::new(
166                    bytemuck::checked::cast_slice(&self.bytes)
167                        .iter()
168                        .map(|e: &f64| e.elem::<E>()),
169                ),
170                // bool is a byte value equal to either 0 or 1
171                DType::Bool(BoolStore::Native) | DType::Bool(BoolStore::U8) => {
172                    Box::new(self.bytes.iter().map(|e| e.elem::<E>()))
173                }
174                DType::Bool(BoolStore::U32) => Box::new(
175                    bytemuck::checked::cast_slice(&self.bytes)
176                        .iter()
177                        .map(|e: &u32| e.elem::<E>()),
178                ),
179                DType::QFloat(scheme) => match scheme {
180                    QuantScheme {
181                        mode: QuantMode::Symmetric,
182                        value:
183                            QuantValue::Q8F
184                            | QuantValue::Q8S
185                            // Represent sub-byte values as i8
186                            | QuantValue::Q4F
187                            | QuantValue::Q4S
188                            | QuantValue::Q2F
189                            | QuantValue::Q2S,
190                        ..
191                    } => {
192                        // Quantized int8 values
193                        let q_bytes = QuantizedBytes {
194                            bytes: self.bytes.clone(),
195                            scheme,
196                            shape: self.shape.clone(),
197                        };
198                        let (values, _) = q_bytes.into_vec_i8();
199
200                        Box::new(
201                            values
202                                .iter()
203                                .map(|e: &i8| e.elem::<E>())
204                                .collect::<Vec<_>>()
205                                .into_iter(),
206                        )
207                    }
208                    QuantScheme {
209                        mode: QuantMode::Symmetric,
210                        value:
211                            QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1,
212                        ..
213                    } => {
214                        unimplemented!("Not yet implemented for iteration");
215                    }
216                    QuantScheme {
217                        mode: QuantMode::Lookup,
218                        ..
219                    } => {
220                        unimplemented!("lookup quantization is not supported for iteration");
221                    }
222                },
223            }
224        }
225    }
226
227    /// Converts the data to the dtype represented by `E`.
228    ///
229    /// # Panics
230    ///
231    /// Panics if storage access fails, the conversion isn't supported, or the stored
232    /// representation or element count is invalid.
233    #[track_caller]
234    pub fn convert<E: Element>(self) -> Self {
235        // TODO: deprecate?
236        self.try_cast_as::<E>()
237            .unwrap_or_else(|err| panic!("Failed to convert TensorData: {err}"))
238    }
239
240    /// Converts the data to `dtype`.
241    ///
242    /// # Panics
243    ///
244    /// Panics if storage access fails, the conversion isn't supported, or the stored
245    /// representation or element count is invalid.
246    #[track_caller]
247    pub fn convert_dtype(self, dtype: DType) -> Self {
248        // TODO: deprecate?
249        self.try_cast(dtype)
250            .unwrap_or_else(|err| panic!("Failed to convert TensorData to {dtype:?}: {err}"))
251    }
252
253    /// Converts the data to the dtype represented by `E`.
254    ///
255    /// By contract, this is equivalent to:
256    /// `data.try_cast(E::dtype())`
257    ///
258    /// # Errors
259    ///
260    /// Returns an error if storage access fails, the conversion isn't supported, or the stored
261    /// representation or element count is invalid.
262    pub fn try_cast_as<E: Element>(self) -> Result<TensorData, DataError> {
263        self.try_cast(E::dtype())
264    }
265
266    /// Converts the data to `dtype`.
267    ///
268    /// # Errors
269    ///
270    /// Returns an error if storage access fails, the conversion isn't supported, or the stored
271    /// representation or element count is invalid.
272    pub fn try_cast(self, dtype: DType) -> Result<TensorData, DataError> {
273        if dtype == self.dtype {
274            Ok(self)
275        } else if dtype.size() == self.dtype.size()
276            && !matches!(
277                self.dtype,
278                DType::Bool(BoolStore::Native) | DType::QFloat(_)
279            )
280            && !matches!(dtype, DType::Bool(BoolStore::Native) | DType::QFloat(_))
281        {
282            self.try_cast_inplace(dtype)
283        } else {
284            self.try_cast_clone(dtype)
285        }
286    }
287
288    // Self-to-Self casts should be stripped before this point.
289    fn try_cast_inplace(self, dtype: DType) -> Result<TensorData, DataError> {
290        // Convert self.dtype to generic parameter:
291        match self.dtype {
292            DType::F64 => self.try_cast_inplace_from::<f64>(dtype),
293            DType::F32 | DType::Flex32 => self.try_cast_inplace_from::<f32>(dtype),
294            DType::F16 => self.try_cast_inplace_from::<f16>(dtype),
295            DType::BF16 => self.try_cast_inplace_from::<bf16>(dtype),
296            DType::I64 => self.try_cast_inplace_from::<i64>(dtype),
297            DType::I32 => self.try_cast_inplace_from::<i32>(dtype),
298            DType::I16 => self.try_cast_inplace_from::<i16>(dtype),
299            DType::I8 => self.try_cast_inplace_from::<i8>(dtype),
300            DType::U64 => self.try_cast_inplace_from::<u64>(dtype),
301            DType::U32 => self.try_cast_inplace_from::<u32>(dtype),
302            DType::U16 => self.try_cast_inplace_from::<u16>(dtype),
303            DType::U8 => self.try_cast_inplace_from::<u8>(dtype),
304            DType::Bool(BoolStore::U8) => self.try_cast_inplace_from::<u8>(dtype),
305            DType::Bool(BoolStore::U32) => self.try_cast_inplace_from::<u32>(dtype),
306            DType::Bool(BoolStore::Native) | DType::QFloat(_) => Err(DataError::DTypeMismatch {
307                expected: dtype,
308                actual: self.dtype,
309            }),
310        }
311    }
312
313    fn try_cast_inplace_from<Current>(self, dtype: DType) -> Result<TensorData, DataError>
314    where
315        Current: Element + AnyBitPattern,
316    {
317        // Convert target dtype to generic parameter.
318        match dtype {
319            DType::F64 => self.try_convert_inplace::<Current, f64>(),
320            DType::F32 | DType::Flex32 => self.try_convert_inplace::<Current, f32>(),
321            DType::F16 => self.try_convert_inplace::<Current, f16>(),
322            DType::BF16 => self.try_convert_inplace::<Current, bf16>(),
323            DType::I64 => self.try_convert_inplace::<Current, i64>(),
324            DType::I32 => self.try_convert_inplace::<Current, i32>(),
325            DType::I16 => self.try_convert_inplace::<Current, i16>(),
326            DType::I8 => self.try_convert_inplace::<Current, i8>(),
327            DType::U64 => self.try_convert_inplace::<Current, u64>(),
328            DType::U32 => self.try_convert_inplace::<Current, u32>(),
329            DType::U16 => self.try_convert_inplace::<Current, u16>(),
330            DType::U8 => self.try_convert_inplace::<Current, u8>(),
331            DType::Bool(BoolStore::U8) => self
332                .try_convert_inplace_bool::<Current, u8>()
333                .map(TensorData::into_bool_u8),
334            DType::Bool(BoolStore::U32) => self
335                .try_convert_inplace_bool::<Current, u32>()
336                .map(TensorData::into_bool_u32),
337            DType::Bool(BoolStore::Native) | DType::QFloat(_) => Err(DataError::DTypeMismatch {
338                expected: dtype,
339                actual: Current::dtype(),
340            }),
341        }
342    }
343
344    fn try_convert_inplace<Current, Target>(self) -> Result<TensorData, DataError>
345    where
346        Current: Element + AnyBitPattern,
347        Target: Element + AnyBitPattern,
348    {
349        self.try_convert_inplace_with::<Current, Target>(|x| x.elem())
350    }
351
352    fn try_convert_inplace_bool<Current, Target>(self) -> Result<TensorData, DataError>
353    where
354        Current: Element + AnyBitPattern,
355        Target: Element + AnyBitPattern,
356    {
357        self.try_convert_inplace_with::<Current, Target>(|x| x.to_bool().elem())
358    }
359
360    fn try_convert_inplace_with<Current, Target>(
361        mut self,
362        transform: impl Fn(&Current) -> Target,
363    ) -> Result<TensorData, DataError>
364    where
365        Current: Element + AnyBitPattern,
366        Target: Element + AnyBitPattern,
367    {
368        let expected = self.num_elements();
369        let values =
370            bytemuck::checked::try_cast_slice_mut::<_, Current>(self.bytes.write(Writer::new())?)
371                .map_err(DataError::InvalidRepresentation)?;
372        let actual = values.len();
373        if actual != expected {
374            return Err(DataError::ElementCountMismatch { expected, actual });
375        }
376
377        for x in values {
378            let t = transform(x);
379            let x = cast_mut::<_, Target>(x);
380            *x = t;
381        }
382
383        self.dtype = Target::dtype();
384
385        Ok(self)
386    }
387
388    fn try_cast_clone(self, dtype: DType) -> Result<TensorData, DataError> {
389        // Convert self.dtype to generic parameter:
390        match self.dtype {
391            DType::F64 => self.try_cast_clone_from::<f64>(dtype),
392            DType::F32 | DType::Flex32 => self.try_cast_clone_from::<f32>(dtype),
393            DType::F16 => self.try_cast_clone_from::<f16>(dtype),
394            DType::BF16 => self.try_cast_clone_from::<bf16>(dtype),
395            DType::I64 => self.try_cast_clone_from::<i64>(dtype),
396            DType::I32 => self.try_cast_clone_from::<i32>(dtype),
397            DType::I16 => self.try_cast_clone_from::<i16>(dtype),
398            DType::I8 => self.try_cast_clone_from::<i8>(dtype),
399            DType::U64 => self.try_cast_clone_from::<u64>(dtype),
400            DType::U32 => self.try_cast_clone_from::<u32>(dtype),
401            DType::U16 => self.try_cast_clone_from::<u16>(dtype),
402            DType::U8 => self.try_cast_clone_from::<u8>(dtype),
403            DType::Bool(BoolStore::Native) => self.try_cast_clone_from::<bool>(dtype),
404            DType::Bool(BoolStore::U8) => self.try_cast_clone_from::<u8>(dtype),
405            DType::Bool(BoolStore::U32) => self.try_cast_clone_from::<u32>(dtype),
406            DType::QFloat(_) => Err(DataError::UnsupportedConversion {
407                to: dtype,
408                from: self.dtype,
409            }),
410        }
411    }
412
413    fn try_cast_clone_from<Current>(self, dtype: DType) -> Result<TensorData, DataError>
414    where
415        Current: Element + CheckedBitPattern,
416    {
417        // Convert target dtype to generic parameter.
418        match dtype {
419            DType::F64 => self.try_convert_clone::<Current, f64>(),
420            DType::F32 | DType::Flex32 => self.try_convert_clone::<Current, f32>(),
421            DType::F16 => self.try_convert_clone::<Current, f16>(),
422            DType::BF16 => self.try_convert_clone::<Current, bf16>(),
423            DType::I64 => self.try_convert_clone::<Current, i64>(),
424            DType::I32 => self.try_convert_clone::<Current, i32>(),
425            DType::I16 => self.try_convert_clone::<Current, i16>(),
426            DType::I8 => self.try_convert_clone::<Current, i8>(),
427            DType::U64 => self.try_convert_clone::<Current, u64>(),
428            DType::U32 => self.try_convert_clone::<Current, u32>(),
429            DType::U16 => self.try_convert_clone::<Current, u16>(),
430            DType::U8 => self.try_convert_clone::<Current, u8>(),
431            DType::Bool(BoolStore::Native) => self.try_convert_clone::<Current, bool>(),
432            DType::Bool(BoolStore::U8) => self
433                .try_convert_clone_bool::<Current, u8>()
434                .map(TensorData::into_bool_u8),
435            DType::Bool(BoolStore::U32) => self
436                .try_convert_clone_bool::<Current, u32>()
437                .map(TensorData::into_bool_u32),
438            DType::QFloat(_) => Err(DataError::UnsupportedConversion {
439                to: dtype,
440                from: self.dtype,
441            }),
442        }
443    }
444
445    fn try_convert_clone<Current, Target>(self) -> Result<TensorData, DataError>
446    where
447        Current: Element + CheckedBitPattern,
448        Target: Element + Zeroable,
449    {
450        self.try_convert_clone_with::<Current, Target>(|x| x.elem())
451    }
452
453    fn try_convert_clone_bool<Current, Target>(self) -> Result<TensorData, DataError>
454    where
455        Current: Element + CheckedBitPattern,
456        Target: Element + Zeroable,
457    {
458        self.try_convert_clone_with::<Current, Target>(|x| x.to_bool().elem())
459    }
460
461    fn try_convert_clone_with<Current, Target>(
462        self,
463        transform: impl Fn(&Current) -> Target,
464    ) -> Result<TensorData, DataError>
465    where
466        Current: Element + CheckedBitPattern,
467        Target: Element + Zeroable,
468    {
469        let expected = self.num_elements();
470        let values =
471            bytemuck::checked::try_cast_slice::<_, Current>(self.bytes.read(Reader::new())?)
472                .map_err(DataError::InvalidRepresentation)?;
473        let actual = values.len();
474        if actual != expected {
475            return Err(DataError::ElementCountMismatch { expected, actual });
476        }
477        let mut out: Vec<Target> = vec![Zeroable::zeroed(); expected];
478
479        for (value, out) in values.iter().zip(&mut out) {
480            *out = transform(value);
481        }
482
483        Ok(TensorData::new(out, self.shape))
484    }
485}
burn_std/data/tensor/conversion.rs

burn_std/data/tensor/
conversion.rs
