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        if E::dtype() == self.dtype {
106            Box::new(bytemuck::checked::cast_slice(&self.bytes).iter().copied())
107        } else {
108            match self.dtype {
109                DType::I8 => Box::new(
110                    bytemuck::checked::cast_slice(&self.bytes)
111                        .iter()
112                        .map(|e: &i8| e.elem::<E>()),
113                ),
114                DType::I16 => Box::new(
115                    bytemuck::checked::cast_slice(&self.bytes)
116                        .iter()
117                        .map(|e: &i16| e.elem::<E>()),
118                ),
119                DType::I32 => Box::new(
120                    bytemuck::checked::cast_slice(&self.bytes)
121                        .iter()
122                        .map(|e: &i32| e.elem::<E>()),
123                ),
124                DType::I64 => Box::new(
125                    bytemuck::checked::cast_slice(&self.bytes)
126                        .iter()
127                        .map(|e: &i64| e.elem::<E>()),
128                ),
129                DType::U8 => Box::new(self.bytes.iter().map(|e| e.elem::<E>())),
130                DType::U16 => Box::new(
131                    bytemuck::checked::cast_slice(&self.bytes)
132                        .iter()
133                        .map(|e: &u16| e.elem::<E>()),
134                ),
135                DType::U32 => Box::new(
136                    bytemuck::checked::cast_slice(&self.bytes)
137                        .iter()
138                        .map(|e: &u32| e.elem::<E>()),
139                ),
140                DType::U64 => Box::new(
141                    bytemuck::checked::cast_slice(&self.bytes)
142                        .iter()
143                        .map(|e: &u64| e.elem::<E>()),
144                ),
145                DType::BF16 => Box::new(
146                    bytemuck::checked::cast_slice(&self.bytes)
147                        .iter()
148                        .map(|e: &bf16| e.elem::<E>()),
149                ),
150                DType::F16 => Box::new(
151                    bytemuck::checked::cast_slice(&self.bytes)
152                        .iter()
153                        .map(|e: &f16| e.elem::<E>()),
154                ),
155                DType::F32 | DType::Flex32 => Box::new(
156                    bytemuck::checked::cast_slice(&self.bytes)
157                        .iter()
158                        .map(|e: &f32| e.elem::<E>()),
159                ),
160                DType::F64 => Box::new(
161                    bytemuck::checked::cast_slice(&self.bytes)
162                        .iter()
163                        .map(|e: &f64| e.elem::<E>()),
164                ),
165                // bool is a byte value equal to either 0 or 1
166                DType::Bool(BoolStore::Native) | DType::Bool(BoolStore::U8) => {
167                    Box::new(self.bytes.iter().map(|e| e.elem::<E>()))
168                }
169                DType::Bool(BoolStore::U32) => Box::new(
170                    bytemuck::checked::cast_slice(&self.bytes)
171                        .iter()
172                        .map(|e: &u32| e.elem::<E>()),
173                ),
174                DType::QFloat(scheme) => match scheme {
175                    QuantScheme {
176                        mode: QuantMode::Symmetric,
177                        value:
178                            QuantValue::Q8F
179                            | QuantValue::Q8S
180                            // Represent sub-byte values as i8
181                            | QuantValue::Q4F
182                            | QuantValue::Q4S
183                            | QuantValue::Q2F
184                            | QuantValue::Q2S,
185                        ..
186                    } => {
187                        // Quantized int8 values
188                        let q_bytes = QuantizedBytes {
189                            bytes: self.bytes.clone(),
190                            scheme,
191                            shape: self.shape.clone(),
192                        };
193                        let (values, _) = q_bytes.into_vec_i8();
194
195                        Box::new(
196                            values
197                                .iter()
198                                .map(|e: &i8| e.elem::<E>())
199                                .collect::<Vec<_>>()
200                                .into_iter(),
201                        )
202                    }
203                    QuantScheme {
204                        mode: QuantMode::Symmetric,
205                        value:
206                            QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1,
207                        ..
208                    } => {
209                        unimplemented!("Not yet implemented for iteration");
210                    }
211                    QuantScheme {
212                        mode: QuantMode::Lookup,
213                        ..
214                    } => {
215                        unimplemented!("lookup quantization is not supported for iteration");
216                    }
217                },
218            }
219        }
220    }
221
222    /// Converts the data to the dtype represented by `E`.
223    ///
224    /// # Panics
225    ///
226    /// Panics if storage access fails, the conversion isn't supported, or the stored
227    /// representation or element count is invalid.
228    #[track_caller]
229    pub fn convert<E: Element>(self) -> Self {
230        // TODO: deprecate?
231        self.try_cast_as::<E>()
232            .unwrap_or_else(|err| panic!("Failed to convert TensorData: {err}"))
233    }
234
235    /// Converts the data to `dtype`.
236    ///
237    /// # Panics
238    ///
239    /// Panics if storage access fails, the conversion isn't supported, or the stored
240    /// representation or element count is invalid.
241    #[track_caller]
242    pub fn convert_dtype(self, dtype: DType) -> Self {
243        // TODO: deprecate?
244        self.try_cast(dtype)
245            .unwrap_or_else(|err| panic!("Failed to convert TensorData to {dtype:?}: {err}"))
246    }
247
248    /// Converts the data to the dtype represented by `E`.
249    ///
250    /// By contract, this is equivalent to:
251    /// `data.try_cast(E::dtype())`
252    ///
253    /// # Errors
254    ///
255    /// Returns an error if storage access fails, the conversion isn't supported, or the stored
256    /// representation or element count is invalid.
257    pub fn try_cast_as<E: Element>(self) -> Result<TensorData, DataError> {
258        self.try_cast(E::dtype())
259    }
260
261    /// Converts the data to `dtype`.
262    ///
263    /// # Errors
264    ///
265    /// Returns an error if storage access fails, the conversion isn't supported, or the stored
266    /// representation or element count is invalid.
267    pub fn try_cast(self, dtype: DType) -> Result<TensorData, DataError> {
268        if dtype == self.dtype {
269            Ok(self)
270        } else if dtype.size() == self.dtype.size()
271            && !matches!(
272                self.dtype,
273                DType::Bool(BoolStore::Native) | DType::QFloat(_)
274            )
275            && !matches!(dtype, DType::Bool(BoolStore::Native) | DType::QFloat(_))
276        {
277            self.try_cast_inplace(dtype)
278        } else {
279            self.try_cast_clone(dtype)
280        }
281    }
282
283    // Self-to-Self casts should be stripped before this point.
284    fn try_cast_inplace(self, dtype: DType) -> Result<TensorData, DataError> {
285        // Convert self.dtype to generic parameter:
286        match self.dtype {
287            DType::F64 => self.try_cast_inplace_from::<f64>(dtype),
288            DType::F32 | DType::Flex32 => self.try_cast_inplace_from::<f32>(dtype),
289            DType::F16 => self.try_cast_inplace_from::<f16>(dtype),
290            DType::BF16 => self.try_cast_inplace_from::<bf16>(dtype),
291            DType::I64 => self.try_cast_inplace_from::<i64>(dtype),
292            DType::I32 => self.try_cast_inplace_from::<i32>(dtype),
293            DType::I16 => self.try_cast_inplace_from::<i16>(dtype),
294            DType::I8 => self.try_cast_inplace_from::<i8>(dtype),
295            DType::U64 => self.try_cast_inplace_from::<u64>(dtype),
296            DType::U32 => self.try_cast_inplace_from::<u32>(dtype),
297            DType::U16 => self.try_cast_inplace_from::<u16>(dtype),
298            DType::U8 => self.try_cast_inplace_from::<u8>(dtype),
299            DType::Bool(BoolStore::U8) => self.try_cast_inplace_from::<u8>(dtype),
300            DType::Bool(BoolStore::U32) => self.try_cast_inplace_from::<u32>(dtype),
301            DType::Bool(BoolStore::Native) | DType::QFloat(_) => Err(DataError::DTypeMismatch {
302                expected: dtype,
303                actual: self.dtype,
304            }),
305        }
306    }
307
308    fn try_cast_inplace_from<Current>(self, dtype: DType) -> Result<TensorData, DataError>
309    where
310        Current: Element + AnyBitPattern,
311    {
312        // Convert target dtype to generic parameter.
313        match dtype {
314            DType::F64 => self.try_convert_inplace::<Current, f64>(),
315            DType::F32 | DType::Flex32 => self.try_convert_inplace::<Current, f32>(),
316            DType::F16 => self.try_convert_inplace::<Current, f16>(),
317            DType::BF16 => self.try_convert_inplace::<Current, bf16>(),
318            DType::I64 => self.try_convert_inplace::<Current, i64>(),
319            DType::I32 => self.try_convert_inplace::<Current, i32>(),
320            DType::I16 => self.try_convert_inplace::<Current, i16>(),
321            DType::I8 => self.try_convert_inplace::<Current, i8>(),
322            DType::U64 => self.try_convert_inplace::<Current, u64>(),
323            DType::U32 => self.try_convert_inplace::<Current, u32>(),
324            DType::U16 => self.try_convert_inplace::<Current, u16>(),
325            DType::U8 => self.try_convert_inplace::<Current, u8>(),
326            DType::Bool(BoolStore::U8) => self
327                .try_convert_inplace_bool::<Current, u8>()
328                .map(TensorData::into_bool_u8),
329            DType::Bool(BoolStore::U32) => self
330                .try_convert_inplace_bool::<Current, u32>()
331                .map(TensorData::into_bool_u32),
332            DType::Bool(BoolStore::Native) | DType::QFloat(_) => Err(DataError::DTypeMismatch {
333                expected: dtype,
334                actual: Current::dtype(),
335            }),
336        }
337    }
338
339    fn try_convert_inplace<Current, Target>(self) -> Result<TensorData, DataError>
340    where
341        Current: Element + AnyBitPattern,
342        Target: Element + AnyBitPattern,
343    {
344        self.try_convert_inplace_with::<Current, Target>(|x| x.elem())
345    }
346
347    fn try_convert_inplace_bool<Current, Target>(self) -> Result<TensorData, DataError>
348    where
349        Current: Element + AnyBitPattern,
350        Target: Element + AnyBitPattern,
351    {
352        self.try_convert_inplace_with::<Current, Target>(|x| x.to_bool().elem())
353    }
354
355    fn try_convert_inplace_with<Current, Target>(
356        mut self,
357        transform: impl Fn(&Current) -> Target,
358    ) -> Result<TensorData, DataError>
359    where
360        Current: Element + AnyBitPattern,
361        Target: Element + AnyBitPattern,
362    {
363        let expected = self.num_elements();
364        let values =
365            bytemuck::checked::try_cast_slice_mut::<_, Current>(self.bytes.write(Writer::new())?)
366                .map_err(DataError::InvalidRepresentation)?;
367        let actual = values.len();
368        if actual != expected {
369            return Err(DataError::ElementCountMismatch { expected, actual });
370        }
371
372        for x in values {
373            let t = transform(x);
374            let x = cast_mut::<_, Target>(x);
375            *x = t;
376        }
377
378        self.dtype = Target::dtype();
379
380        Ok(self)
381    }
382
383    fn try_cast_clone(self, dtype: DType) -> Result<TensorData, DataError> {
384        // Convert self.dtype to generic parameter:
385        match self.dtype {
386            DType::F64 => self.try_cast_clone_from::<f64>(dtype),
387            DType::F32 | DType::Flex32 => self.try_cast_clone_from::<f32>(dtype),
388            DType::F16 => self.try_cast_clone_from::<f16>(dtype),
389            DType::BF16 => self.try_cast_clone_from::<bf16>(dtype),
390            DType::I64 => self.try_cast_clone_from::<i64>(dtype),
391            DType::I32 => self.try_cast_clone_from::<i32>(dtype),
392            DType::I16 => self.try_cast_clone_from::<i16>(dtype),
393            DType::I8 => self.try_cast_clone_from::<i8>(dtype),
394            DType::U64 => self.try_cast_clone_from::<u64>(dtype),
395            DType::U32 => self.try_cast_clone_from::<u32>(dtype),
396            DType::U16 => self.try_cast_clone_from::<u16>(dtype),
397            DType::U8 => self.try_cast_clone_from::<u8>(dtype),
398            DType::Bool(BoolStore::Native) => self.try_cast_clone_from::<bool>(dtype),
399            DType::Bool(BoolStore::U8) => self.try_cast_clone_from::<u8>(dtype),
400            DType::Bool(BoolStore::U32) => self.try_cast_clone_from::<u32>(dtype),
401            DType::QFloat(_) => Err(DataError::UnsupportedConversion {
402                to: dtype,
403                from: self.dtype,
404            }),
405        }
406    }
407
408    fn try_cast_clone_from<Current>(self, dtype: DType) -> Result<TensorData, DataError>
409    where
410        Current: Element + CheckedBitPattern,
411    {
412        // Convert target dtype to generic parameter.
413        match dtype {
414            DType::F64 => self.try_convert_clone::<Current, f64>(),
415            DType::F32 | DType::Flex32 => self.try_convert_clone::<Current, f32>(),
416            DType::F16 => self.try_convert_clone::<Current, f16>(),
417            DType::BF16 => self.try_convert_clone::<Current, bf16>(),
418            DType::I64 => self.try_convert_clone::<Current, i64>(),
419            DType::I32 => self.try_convert_clone::<Current, i32>(),
420            DType::I16 => self.try_convert_clone::<Current, i16>(),
421            DType::I8 => self.try_convert_clone::<Current, i8>(),
422            DType::U64 => self.try_convert_clone::<Current, u64>(),
423            DType::U32 => self.try_convert_clone::<Current, u32>(),
424            DType::U16 => self.try_convert_clone::<Current, u16>(),
425            DType::U8 => self.try_convert_clone::<Current, u8>(),
426            DType::Bool(BoolStore::Native) => self.try_convert_clone::<Current, bool>(),
427            DType::Bool(BoolStore::U8) => self
428                .try_convert_clone_bool::<Current, u8>()
429                .map(TensorData::into_bool_u8),
430            DType::Bool(BoolStore::U32) => self
431                .try_convert_clone_bool::<Current, u32>()
432                .map(TensorData::into_bool_u32),
433            DType::QFloat(_) => Err(DataError::UnsupportedConversion {
434                to: dtype,
435                from: self.dtype,
436            }),
437        }
438    }
439
440    fn try_convert_clone<Current, Target>(self) -> Result<TensorData, DataError>
441    where
442        Current: Element + CheckedBitPattern,
443        Target: Element + Zeroable,
444    {
445        self.try_convert_clone_with::<Current, Target>(|x| x.elem())
446    }
447
448    fn try_convert_clone_bool<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.to_bool().elem())
454    }
455
456    fn try_convert_clone_with<Current, Target>(
457        self,
458        transform: impl Fn(&Current) -> Target,
459    ) -> Result<TensorData, DataError>
460    where
461        Current: Element + CheckedBitPattern,
462        Target: Element + Zeroable,
463    {
464        let expected = self.num_elements();
465        let values =
466            bytemuck::checked::try_cast_slice::<_, Current>(self.bytes.read(Reader::new())?)
467                .map_err(DataError::InvalidRepresentation)?;
468        let actual = values.len();
469        if actual != expected {
470            return Err(DataError::ElementCountMismatch { expected, actual });
471        }
472        let mut out: Vec<Target> = vec![Zeroable::zeroed(); expected];
473
474        for (value, out) in values.iter().zip(&mut out) {
475            *out = transform(value);
476        }
477
478        Ok(TensorData::new(out, self.shape))
479    }
480}
burn_std/data/tensor/conversion.rs

burn_std/data/tensor/
conversion.rs
