burn_std/tensor/
dtype.rs

1//! Tensor data type.
2
3use serde::{Deserialize, Serialize};
4
5use crate::tensor::quantization::{QuantScheme, QuantStore, QuantValue};
6use crate::{bf16, f16};
7
8#[allow(missing_docs)]
9#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Serialize, Deserialize)]
10pub enum DType {
11    F64,
12    F32,
13    Flex32,
14    F16,
15    BF16,
16    I64,
17    I32,
18    I16,
19    I8,
20    U64,
21    U32,
22    U16,
23    U8,
24    Bool(BoolStore),
25    QFloat(QuantScheme),
26}
27
28// Conversions between `DType` and cubecl's `ElemType` / `StorageType` are
29// intentionally not implemented here so that `burn-std` does not depend on
30// `cubecl`. Backend code that needs them should call the named functions in
31// `burn_backend::cubecl` (e.g. `elem_type_to_dtype`, `dtype_to_elem_type`,
32// `dtype_to_storage_type`).
33
34impl DType {
35    /// Returns the size of a type in bytes.
36    pub const fn size(&self) -> usize {
37        match self {
38            DType::F64 => core::mem::size_of::<f64>(),
39            DType::F32 => core::mem::size_of::<f32>(),
40            DType::Flex32 => core::mem::size_of::<f32>(),
41            DType::F16 => core::mem::size_of::<f16>(),
42            DType::BF16 => core::mem::size_of::<bf16>(),
43            DType::I64 => core::mem::size_of::<i64>(),
44            DType::I32 => core::mem::size_of::<i32>(),
45            DType::I16 => core::mem::size_of::<i16>(),
46            DType::I8 => core::mem::size_of::<i8>(),
47            DType::U64 => core::mem::size_of::<u64>(),
48            DType::U32 => core::mem::size_of::<u32>(),
49            DType::U16 => core::mem::size_of::<u16>(),
50            DType::U8 => core::mem::size_of::<u8>(),
51            DType::Bool(store) => match store {
52                BoolStore::Native => core::mem::size_of::<bool>(),
53                BoolStore::U8 => core::mem::size_of::<u8>(),
54                BoolStore::U32 => core::mem::size_of::<u32>(),
55            },
56            DType::QFloat(scheme) => match scheme.store {
57                QuantStore::Native => match scheme.value {
58                    QuantValue::Q8F | QuantValue::Q8S => core::mem::size_of::<i8>(),
59                    // e2m1 native is automatically packed by the kernels, so the actual storage is
60                    // 8 bits wide.
61                    QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1 => {
62                        core::mem::size_of::<u8>()
63                    }
64                    QuantValue::Q4F | QuantValue::Q4S | QuantValue::Q2F | QuantValue::Q2S => {
65                        // Sub-byte values have fractional size
66                        0
67                    }
68                },
69                QuantStore::PackedU32(_) => core::mem::size_of::<u32>(),
70                QuantStore::PackedNative(_) => match scheme.value {
71                    QuantValue::E2M1 => core::mem::size_of::<u8>(),
72                    _ => 0,
73                },
74            },
75        }
76    }
77    /// Returns true if the data type is a floating point type.
78    pub fn is_float(&self) -> bool {
79        matches!(
80            self,
81            DType::F64 | DType::F32 | DType::Flex32 | DType::F16 | DType::BF16
82        )
83    }
84    /// Returns true if the data type is a signed integer type.
85    pub fn is_int(&self) -> bool {
86        matches!(self, DType::I64 | DType::I32 | DType::I16 | DType::I8)
87    }
88    /// Returns true if the data type is an unsigned integer type.
89    pub fn is_uint(&self) -> bool {
90        matches!(self, DType::U64 | DType::U32 | DType::U16 | DType::U8)
91    }
92
93    /// Returns true if the data type is a boolean type
94    pub fn is_bool(&self) -> bool {
95        matches!(self, DType::Bool(_))
96    }
97
98    /// Returns float precision info if this is a float dtype, `None` otherwise.
99    ///
100    /// Analogous to `torch.finfo(dtype)` or `numpy.finfo(dtype)`.
101    pub const fn finfo(&self) -> Option<FloatInfo> {
102        match self {
103            DType::F64 => Some(FloatDType::F64.finfo()),
104            DType::F32 => Some(FloatDType::F32.finfo()),
105            DType::Flex32 => Some(FloatDType::Flex32.finfo()),
106            DType::F16 => Some(FloatDType::F16.finfo()),
107            DType::BF16 => Some(FloatDType::BF16.finfo()),
108            _ => None,
109        }
110    }
111
112    /// Returns the data type name.
113    pub fn name(&self) -> &'static str {
114        match self {
115            DType::F64 => "f64",
116            DType::F32 => "f32",
117            DType::Flex32 => "flex32",
118            DType::F16 => "f16",
119            DType::BF16 => "bf16",
120            DType::I64 => "i64",
121            DType::I32 => "i32",
122            DType::I16 => "i16",
123            DType::I8 => "i8",
124            DType::U64 => "u64",
125            DType::U32 => "u32",
126            DType::U16 => "u16",
127            DType::U8 => "u8",
128            DType::Bool(store) => match store {
129                BoolStore::Native => "bool",
130                BoolStore::U8 => "bool(u8)",
131                BoolStore::U32 => "bool(u32)",
132            },
133            DType::QFloat(_) => "qfloat",
134        }
135    }
136}
137
138#[allow(missing_docs)]
139#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
140pub enum FloatDType {
141    F64,
142    F32,
143    Flex32,
144    F16,
145    BF16,
146}
147
148/// Numerical precision properties for a floating-point dtype.
149///
150/// Equivalent to NumPy's `finfo` / PyTorch's `torch.finfo`. All values are
151/// widened to `f64` so they can be inspected without knowing the concrete
152/// element type at compile time.
153#[derive(Debug, Clone, Copy, PartialEq)]
154pub struct FloatInfo {
155    /// Machine epsilon: smallest value such that `1.0 + epsilon != 1.0`.
156    pub epsilon: f64,
157    /// Largest representable finite value.
158    pub max: f64,
159    /// Most negative representable finite value.
160    pub min: f64,
161    /// Smallest positive normal value.
162    pub min_positive: f64,
163}
164
165impl FloatDType {
166    /// Returns numerical precision properties for this float dtype.
167    ///
168    /// Analogous to `torch.finfo(dtype)` or `numpy.finfo(dtype)`.
169    pub const fn finfo(self) -> FloatInfo {
170        match self {
171            FloatDType::F64 => FloatInfo {
172                epsilon: f64::EPSILON,
173                max: f64::MAX,
174                min: f64::MIN,
175                min_positive: f64::MIN_POSITIVE, // ~2.225e-308
176            },
177            FloatDType::F32 => FloatInfo {
178                epsilon: f32::EPSILON as f64,
179                max: f32::MAX as f64,
180                min: f32::MIN as f64,
181                min_positive: f32::MIN_POSITIVE as f64, // ~1.175e-38
182            },
183            // Flex32 stores as f32 but computes at reduced (f16-like) precision.
184            // Use f16 precision limits so stability code stays safe.
185            FloatDType::Flex32 => FloatInfo {
186                epsilon: f16::EPSILON.to_f64_const(),
187                max: f16::MAX.to_f64_const(),
188                min: f16::MIN.to_f64_const(),
189                min_positive: f16::MIN_POSITIVE.to_f64_const(), // ~6.104e-5
190            },
191            FloatDType::F16 => FloatInfo {
192                epsilon: f16::EPSILON.to_f64_const(),
193                max: f16::MAX.to_f64_const(),
194                min: f16::MIN.to_f64_const(),
195                min_positive: f16::MIN_POSITIVE.to_f64_const(), // ~6.104e-5
196            },
197            FloatDType::BF16 => FloatInfo {
198                epsilon: bf16::EPSILON.to_f64_const(),
199                max: bf16::MAX.to_f64_const(),
200                min: bf16::MIN.to_f64_const(),
201                min_positive: bf16::MIN_POSITIVE.to_f64_const(), // ~1.175e-38
202            },
203        }
204    }
205}
206
207impl From<DType> for FloatDType {
208    fn from(value: DType) -> Self {
209        match value {
210            DType::F64 => FloatDType::F64,
211            DType::F32 => FloatDType::F32,
212            DType::Flex32 => FloatDType::Flex32,
213            DType::F16 => FloatDType::F16,
214            DType::BF16 => FloatDType::BF16,
215            _ => panic!("Expected float data type, got {value:?}"),
216        }
217    }
218}
219
220impl From<FloatDType> for DType {
221    fn from(value: FloatDType) -> Self {
222        match value {
223            FloatDType::F64 => DType::F64,
224            FloatDType::F32 => DType::F32,
225            FloatDType::Flex32 => DType::Flex32,
226            FloatDType::F16 => DType::F16,
227            FloatDType::BF16 => DType::BF16,
228        }
229    }
230}
231
232#[allow(missing_docs)]
233#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
234pub enum IntDType {
235    I64,
236    I32,
237    I16,
238    I8,
239    U64,
240    U32,
241    U16,
242    U8,
243}
244
245impl From<DType> for IntDType {
246    fn from(value: DType) -> Self {
247        match value {
248            DType::I64 => IntDType::I64,
249            DType::I32 => IntDType::I32,
250            DType::I16 => IntDType::I16,
251            DType::I8 => IntDType::I8,
252            DType::U64 => IntDType::U64,
253            DType::U32 => IntDType::U32,
254            DType::U16 => IntDType::U16,
255            DType::U8 => IntDType::U8,
256            _ => panic!("Expected int data type, got {value:?}"),
257        }
258    }
259}
260
261impl From<IntDType> for DType {
262    fn from(value: IntDType) -> Self {
263        match value {
264            IntDType::I64 => DType::I64,
265            IntDType::I32 => DType::I32,
266            IntDType::I16 => DType::I16,
267            IntDType::I8 => DType::I8,
268            IntDType::U64 => DType::U64,
269            IntDType::U32 => DType::U32,
270            IntDType::U16 => DType::U16,
271            IntDType::U8 => DType::U8,
272        }
273    }
274}
275
276#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Serialize, Deserialize)]
277/// Data type used to store boolean values.
278pub enum BoolStore {
279    /// Stored as native boolean type (e.g. `bool`).
280    Native,
281    /// Stored as 8-bit unsigned integer.
282    U8,
283    /// Stored as 32-bit unsigned integer.
284    U32,
285}
286
287/// Boolean dtype.
288///
289/// This is currently an alias to [`BoolStore`], since it only varies by the storage representation.
290pub type BoolDType = BoolStore;
291
292#[allow(deprecated)]
293impl From<DType> for BoolDType {
294    fn from(value: DType) -> Self {
295        match value {
296            DType::Bool(store) => match store {
297                BoolStore::Native => BoolDType::Native,
298                BoolStore::U8 => BoolDType::U8,
299                BoolStore::U32 => BoolDType::U32,
300            },
301            // For compat BoolElem associated type
302            DType::U8 => BoolDType::U8,
303            DType::U32 => BoolDType::U32,
304            _ => panic!("Expected bool data type, got {value:?}"),
305        }
306    }
307}
308
309impl From<BoolDType> for DType {
310    fn from(value: BoolDType) -> Self {
311        match value {
312            BoolDType::Native => DType::Bool(BoolStore::Native),
313            BoolDType::U8 => DType::Bool(BoolStore::U8),
314            BoolDType::U32 => DType::Bool(BoolStore::U32),
315        }
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322
323    #[test]
324    fn finfo_f32() {
325        let info = FloatDType::F32.finfo();
326        assert_eq!(info.epsilon, f32::EPSILON as f64);
327        assert_eq!(info.max, f32::MAX as f64);
328        assert_eq!(info.min, f32::MIN as f64);
329        assert_eq!(info.min_positive, f32::MIN_POSITIVE as f64);
330    }
331
332    #[test]
333    fn finfo_f64() {
334        let info = FloatDType::F64.finfo();
335        assert_eq!(info.epsilon, f64::EPSILON);
336        assert_eq!(info.max, f64::MAX);
337        assert_eq!(info.min, f64::MIN);
338        assert_eq!(info.min_positive, f64::MIN_POSITIVE);
339    }
340
341    #[test]
342    fn finfo_f16() {
343        let info = FloatDType::F16.finfo();
344        assert_eq!(info.epsilon, f16::EPSILON.to_f64_const());
345        assert!(info.epsilon > 0.0);
346        assert!(info.min_positive > 0.0);
347        // f16 epsilon is much larger than f32
348        assert!(info.epsilon > FloatDType::F32.finfo().epsilon);
349    }
350
351    #[test]
352    fn finfo_bf16() {
353        let info = FloatDType::BF16.finfo();
354        assert_eq!(info.epsilon, bf16::EPSILON.to_f64_const());
355        assert!(info.epsilon > 0.0);
356        assert!(info.min_positive > 0.0);
357        // bf16 epsilon is larger than f32 (fewer mantissa bits)
358        assert!(info.epsilon > FloatDType::F32.finfo().epsilon);
359    }
360
361    #[test]
362    fn finfo_flex32_uses_f16_limits() {
363        let flex = FloatDType::Flex32.finfo();
364        let f16_info = FloatDType::F16.finfo();
365        assert_eq!(flex.epsilon, f16_info.epsilon);
366        assert_eq!(flex.min_positive, f16_info.min_positive);
367    }
368
369    #[test]
370    fn dtype_finfo_delegates_to_float_dtype() {
371        assert_eq!(DType::F32.finfo(), Some(FloatDType::F32.finfo()));
372        assert_eq!(DType::F64.finfo(), Some(FloatDType::F64.finfo()));
373        assert_eq!(DType::F16.finfo(), Some(FloatDType::F16.finfo()));
374        assert_eq!(DType::BF16.finfo(), Some(FloatDType::BF16.finfo()));
375        assert_eq!(DType::Flex32.finfo(), Some(FloatDType::Flex32.finfo()));
376    }
377
378    #[test]
379    fn dtype_finfo_returns_none_for_non_float() {
380        assert!(DType::I32.finfo().is_none());
381        assert!(DType::U8.finfo().is_none());
382        assert!(DType::Bool(BoolStore::Native).finfo().is_none());
383    }
384
385    #[test]
386    fn finfo_invariants() {
387        for dtype in [
388            FloatDType::F64,
389            FloatDType::F32,
390            FloatDType::F16,
391            FloatDType::BF16,
392            FloatDType::Flex32,
393        ] {
394            let info = dtype.finfo();
395            assert!(info.epsilon > 0.0, "{dtype:?}: epsilon must be positive");
396            assert!(
397                info.min_positive > 0.0,
398                "{dtype:?}: min_positive must be positive"
399            );
400            assert!(info.max > 0.0, "{dtype:?}: max must be positive");
401            assert!(info.min < 0.0, "{dtype:?}: min must be negative");
402            assert!(
403                info.max > info.min_positive,
404                "{dtype:?}: max > min_positive"
405            );
406        }
407    }
408}
burn_std/tensor/dtype.rs

burn_std/tensor/
dtype.rs
