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

burn_std/data/tensor/
conversion.rs
