kime_tensor/
dtype.rs

1//! Element types as they appear in checkpoints.
2
3use std::fmt;
4
5/// The element types kime reads from checkpoints. Weights are floats. The integer types exist
6/// because safetensors files may carry them, and a loader that knows their size can check a file
7/// it will not use.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9pub enum DType {
10    /// IEEE binary32.
11    F32,
12    /// IEEE binary16.
13    F16,
14    /// bfloat16.
15    BF16,
16    /// Signed 64 bit integers.
17    I64,
18    /// Signed 32 bit integers.
19    I32,
20    /// Signed 8 bit integers.
21    I8,
22    /// Unsigned 8 bit integers.
23    U8,
24    /// Booleans, one byte each.
25    Bool,
26}
27
28impl DType {
29    /// Bytes per element.
30    #[must_use]
31    pub const fn size(self) -> usize {
32        match self {
33            Self::F32 | Self::I32 => 4,
34            Self::F16 | Self::BF16 => 2,
35            Self::I64 => 8,
36            Self::I8 | Self::U8 | Self::Bool => 1,
37        }
38    }
39
40    /// The name safetensors uses, which kime also uses in `.kime` headers.
41    #[must_use]
42    pub const fn name(self) -> &'static str {
43        match self {
44            Self::F32 => "F32",
45            Self::F16 => "F16",
46            Self::BF16 => "BF16",
47            Self::I64 => "I64",
48            Self::I32 => "I32",
49            Self::I8 => "I8",
50            Self::U8 => "U8",
51            Self::Bool => "BOOL",
52        }
53    }
54
55    /// The inverse of [`DType::name`].
56    #[must_use]
57    pub fn from_name(name: &str) -> Option<Self> {
58        Some(match name {
59            "F32" => Self::F32,
60            "F16" => Self::F16,
61            "BF16" => Self::BF16,
62            "I64" => Self::I64,
63            "I32" => Self::I32,
64            "I8" => Self::I8,
65            "U8" => Self::U8,
66            "BOOL" => Self::Bool,
67            _ => return None,
68        })
69    }
70
71    /// Whether this is one of the float types weights come in.
72    #[must_use]
73    pub const fn is_float(self) -> bool {
74        matches!(self, Self::F32 | Self::F16 | Self::BF16)
75    }
76
77    /// Reads element `i` of little endian `bytes` as f32. Panics if `bytes` is too short, and
78    /// returns NaN for the integer types, which no weight uses.
79    #[must_use]
80    pub fn read_f32(self, bytes: &[u8], i: usize) -> f32 {
81        let s = self.size();
82        let b = &bytes[i * s..(i + 1) * s];
83        match self {
84            Self::F32 => f32::from_le_bytes([b[0], b[1], b[2], b[3]]),
85            Self::F16 => half::f16::from_le_bytes([b[0], b[1]]).to_f32(),
86            Self::BF16 => half::bf16::from_le_bytes([b[0], b[1]]).to_f32(),
87            _ => f32::NAN,
88        }
89    }
90}
91
92impl fmt::Display for DType {
93    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
94        f.write_str(self.name())
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101
102    #[test]
103    fn names_round_trip() {
104        for d in [
105            DType::F32,
106            DType::F16,
107            DType::BF16,
108            DType::I64,
109            DType::I32,
110            DType::I8,
111            DType::U8,
112            DType::Bool,
113        ] {
114            assert_eq!(DType::from_name(d.name()), Some(d));
115        }
116        assert_eq!(DType::from_name("F64"), None);
117    }
118
119    #[test]
120    fn reads_halves() {
121        let one = half::f16::from_f32(1.5).to_le_bytes();
122        assert_eq!(DType::F16.read_f32(&one, 0).to_bits(), 1.5f32.to_bits());
123        let two = half::bf16::from_f32(-2.0).to_le_bytes();
124        assert_eq!(DType::BF16.read_f32(&two, 0).to_bits(), (-2.0f32).to_bits());
125    }
126}
kime_tensor/dtype.rs

kime_tensor/
dtype.rs
