mlx_rs/
dtype.rs

1use half::{bf16, f16};
2use mlx_internal_macros::generate_test_cases;
3use strum::EnumIter;
4
5use crate::error::InexactDtypeError;
6
7generate_test_cases! {
8    /// Array element type
9    #[derive(
10        Debug,
11        Clone,
12        Copy,
13        PartialEq,
14        Eq,
15        num_enum::IntoPrimitive,
16        num_enum::TryFromPrimitive,
17        EnumIter,
18        Hash,
19    )]
20    #[repr(u32)]
21    pub enum Dtype {
22        /// bool
23        Bool = mlx_sys::mlx_dtype__MLX_BOOL,
24
25        /// u8
26        Uint8 = mlx_sys::mlx_dtype__MLX_UINT8,
27
28        /// u16
29        Uint16 = mlx_sys::mlx_dtype__MLX_UINT16,
30
31        /// u32
32        Uint32 = mlx_sys::mlx_dtype__MLX_UINT32,
33
34        /// u64
35        Uint64 = mlx_sys::mlx_dtype__MLX_UINT64,
36
37        /// i8
38        Int8 = mlx_sys::mlx_dtype__MLX_INT8,
39
40        /// i16
41        Int16 = mlx_sys::mlx_dtype__MLX_INT16,
42
43        /// i32
44        Int32 = mlx_sys::mlx_dtype__MLX_INT32,
45
46        /// i64
47        Int64 = mlx_sys::mlx_dtype__MLX_INT64,
48
49        /// f16
50        Float16 = mlx_sys::mlx_dtype__MLX_FLOAT16,
51
52        /// f32
53        Float32 = mlx_sys::mlx_dtype__MLX_FLOAT32,
54
55        /// f64
56        Float64 = mlx_sys::mlx_dtype__MLX_FLOAT64,
57
58        /// bfloat16
59        Bfloat16 = mlx_sys::mlx_dtype__MLX_BFLOAT16,
60
61        /// complex64
62        Complex64 = mlx_sys::mlx_dtype__MLX_COMPLEX64,
63    }
64}
65
66impl Dtype {
67    /// Returns `true` if the data type is complex.
68    pub fn is_complex(&self) -> bool {
69        matches!(self, Dtype::Complex64)
70    }
71
72    /// Returns `true` if the data type is floating point.
73    pub fn is_float(&self) -> bool {
74        matches!(
75            self,
76            Dtype::Float16 | Dtype::Float32 | Dtype::Float64 | Dtype::Bfloat16
77        )
78    }
79
80    /// Returns `true` if the data type is one of `f16`, `f32`, `f64`, `bfloat16`, or `complex64`.
81    pub fn is_inexact(&self) -> bool {
82        matches!(
83            self,
84            Dtype::Float16 | Dtype::Float32 | Dtype::Float64 | Dtype::Complex64 | Dtype::Bfloat16
85        )
86    }
87
88    /// Returns the promotion type of two data types.
89    pub fn from_promoting_types(a: Dtype, b: Dtype) -> Self {
90        a.promote_with(b)
91    }
92
93    /// Minimum value of the float point types. Returns `Err(_)` if the type is not
94    /// float point
95    pub fn finfo_min(&self) -> Result<f64, InexactDtypeError> {
96        match self {
97            Dtype::Float16 => Ok(f16::MIN.to_f64_const()),
98            Dtype::Float32 => Ok(f32::MIN as f64),
99            Dtype::Complex64 => Ok(f32::MIN as f64),
100            Dtype::Bfloat16 => Ok(bf16::MIN.to_f64_const()),
101            _ => Err(InexactDtypeError(*self)),
102        }
103    }
104
105    /// Maximum value of the float point types. Returns `Err(_)` if the type is not
106    /// float point
107    pub fn finfo_max(&self) -> Result<f64, InexactDtypeError> {
108        match self {
109            Dtype::Float16 => Ok(f16::MAX.to_f64_const()),
110            Dtype::Float32 => Ok(f32::MAX as f64),
111            Dtype::Complex64 => Ok(f32::MAX as f64),
112            Dtype::Bfloat16 => Ok(bf16::MAX.to_f64_const()),
113            _ => Err(InexactDtypeError(*self)),
114        }
115    }
116}
117
118pub(crate) trait TypePromotion {
119    fn promote_with(self, other: Self) -> Self;
120}
121
122impl TypePromotion for Dtype {
123    fn promote_with(self, other: Self) -> Self {
124        use crate::dtype::Dtype::*;
125        match (self, other) {
126            // Boolean promotions
127            (Bool, Bool) => Bool,
128            (Bool, _) | (_, Bool) => {
129                if self == Bool {
130                    other
131                } else {
132                    self
133                }
134            }
135
136            // Uint8 promotions
137            (Uint8, Uint8) => Uint8,
138            (Uint8, Uint16) | (Uint16, Uint8) => Uint16,
139            (Uint8, Uint32) | (Uint32, Uint8) => Uint32,
140            (Uint8, Uint64) | (Uint64, Uint8) => Uint64,
141            (Uint8, Int8) | (Int8, Uint8) => Int16,
142            (Uint8, Int16) | (Int16, Uint8) => Int16,
143            (Uint8, Int32) | (Int32, Uint8) => Int32,
144            (Uint8, Int64) | (Int64, Uint8) => Int64,
145
146            // Uint16 promotions
147            (Uint16, Uint16) => Uint16,
148            (Uint16, Uint32) | (Uint32, Uint16) => Uint32,
149            (Uint16, Uint64) | (Uint64, Uint16) => Uint64,
150            (Uint16, Int8) | (Int8, Uint16) => Int32,
151            (Uint16, Int16) | (Int16, Uint16) => Int32,
152            (Uint16, Int32) | (Int32, Uint16) => Int32,
153            (Uint16, Int64) | (Int64, Uint16) => Int64,
154
155            // Uint32 promotions
156            (Uint32, Uint32) => Uint32,
157            (Uint32, Uint64) | (Uint64, Uint32) => Uint64,
158            (Uint32, Int8) | (Int8, Uint32) => Int64,
159            (Uint32, Int16) | (Int16, Uint32) => Int64,
160            (Uint32, Int32) | (Int32, Uint32) => Int64,
161            (Uint32, Int64) | (Int64, Uint32) => Int64,
162
163            // Uint64 promotions
164            (Uint64, Uint64) => Uint64,
165            (Uint64, Int8) | (Int8, Uint64) => Float32,
166            (Uint64, Int16) | (Int16, Uint64) => Float32,
167            (Uint64, Int32) | (Int32, Uint64) => Float32,
168            (Uint64, Int64) | (Int64, Uint64) => Float32,
169
170            // Int8 promotions
171            (Int8, Int8) => Int8,
172            (Int8, Int16) | (Int16, Int8) => Int16,
173            (Int8, Int32) | (Int32, Int8) => Int32,
174            (Int8, Int64) | (Int64, Int8) => Int64,
175
176            // Int16 promotions
177            (Int16, Int16) => Int16,
178            (Int16, Int32) | (Int32, Int16) => Int32,
179            (Int16, Int64) | (Int64, Int16) => Int64,
180
181            // Int32 promotions
182            (Int32, Int32) => Int32,
183            (Int32, Int64) | (Int64, Int32) => Int64,
184
185            // Int64 promotions
186            (Int64, Int64) => Int64,
187
188            // Float16 promotions
189            (Float16, Bfloat16) | (Bfloat16, Float16) => Float32,
190
191            // Complex type
192            (Complex64, _) | (_, Complex64) => Complex64,
193
194            // Float64 promotions
195            (Float64, _) | (_, Float64) => Float64,
196
197            // Float32 promotions
198            (Float32, _) | (_, Float32) => Float32,
199
200            // Float16 promotions
201            (Float16, _) | (_, Float16) => Float16,
202
203            // Bfloat16 promotions
204            (Bfloat16, _) | (_, Bfloat16) => Bfloat16,
205        }
206    }
207}
208
209cfg_safetensors! {
210    impl TryFrom<safetensors::tensor::Dtype> for Dtype {
211        type Error = crate::error::ConversionError;
212
213        fn try_from(value: safetensors::tensor::Dtype) -> Result<Self, Self::Error> {
214            let out = match value {
215                safetensors::Dtype::BOOL => Dtype::Bool,
216                safetensors::Dtype::U8 => Dtype::Uint8,
217                safetensors::Dtype::I8 => Dtype::Int8,
218                safetensors::Dtype::F8_E5M2 => return Err(crate::error::ConversionError::SafeTensorDtype(value)),
219                safetensors::Dtype::F8_E4M3 => return Err(crate::error::ConversionError::SafeTensorDtype(value)),
220                safetensors::Dtype::I16 => Dtype::Int16,
221                safetensors::Dtype::U16 => Dtype::Uint16,
222                safetensors::Dtype::F16 => Dtype::Float16,
223                safetensors::Dtype::BF16 => Dtype::Bfloat16,
224                safetensors::Dtype::I32 => Dtype::Int32,
225                safetensors::Dtype::U32 => Dtype::Uint32,
226                safetensors::Dtype::F32 => Dtype::Float32,
227                safetensors::Dtype::F64 => Dtype::Float64,
228                safetensors::Dtype::I64 => Dtype::Int64,
229                safetensors::Dtype::U64 => Dtype::Uint64,
230                _ => return Err(crate::error::ConversionError::SafeTensorDtype(value)),
231            };
232            Ok(out)
233        }
234    }
235
236    impl TryFrom<Dtype> for safetensors::tensor::Dtype {
237        type Error = crate::error::ConversionError;
238
239        fn try_from(value: Dtype) -> Result<Self, Self::Error> {
240            let out = match value {
241                Dtype::Bool => safetensors::Dtype::BOOL,
242                Dtype::Uint8 => safetensors::Dtype::U8,
243                Dtype::Int8 => safetensors::Dtype::I8,
244                Dtype::Int16 => safetensors::Dtype::I16,
245                Dtype::Uint16 => safetensors::Dtype::U16,
246                Dtype::Float16 => safetensors::Dtype::F16,
247                Dtype::Bfloat16 => safetensors::Dtype::BF16,
248                Dtype::Int32 => safetensors::Dtype::I32,
249                Dtype::Uint32 => safetensors::Dtype::U32,
250                Dtype::Float32 => safetensors::Dtype::F32,
251                Dtype::Float64 => safetensors::Dtype::F64,
252                Dtype::Int64 => safetensors::Dtype::I64,
253                Dtype::Uint64 => safetensors::Dtype::U64,
254                Dtype::Complex64 => return Err(crate::error::ConversionError::MlxDtype(value)),
255            };
256            Ok(out)
257        }
258    }
259}
mlx_rs/dtype.rs

mlx_rs/
dtype.rs
