From fdd2155607428ed25610617344c6742fb2a4be55 Mon Sep 17 00:00:00 2001 From: kerthcet Date: Tue, 25 Aug 2026 00:53:23 +0100 Subject: [PATCH] Add Dtype Signed-off-by: kerthcet --- crates/mlxcore/src/array.rs | 404 +++++++++++++++++++++++++++++++++-- crates/mlxcore/src/dtype.rs | 190 ++++++++++++++-- crates/mlxcore/src/lib.rs | 2 +- crates/mlxcore/src/random.rs | 4 +- 4 files changed, 566 insertions(+), 34 deletions(-) diff --git a/crates/mlxcore/src/array.rs b/crates/mlxcore/src/array.rs index ddd883e..748a28e 100644 --- a/crates/mlxcore/src/array.rs +++ b/crates/mlxcore/src/array.rs @@ -5,7 +5,7 @@ use std::fmt; use mlxcore_sys as sys; -use crate::dtype::ArrayElement; +use crate::dtype::{ArrayElement, Dtype}; use crate::error::{self, Result}; use crate::ffi::as_ffi_ptr; use crate::stream::Stream; @@ -56,8 +56,9 @@ impl Array { let shape_ptr = as_ffi_ptr(shape); // SAFETY: pointers/len describe valid slices (or null/0) for the // duration of the call; mlx copies the data into its own buffer. - let handle = - unsafe { sys::mlx_array_new_data(data_ptr, shape_ptr, shape.len() as i32, T::DTYPE) }; + let handle = unsafe { + sys::mlx_array_new_data(data_ptr, shape_ptr, shape.len() as i32, T::DTYPE.as_raw()) + }; unsafe { Self::from_raw(handle) } } @@ -98,7 +99,7 @@ impl Array { shape_ptr, shape.len(), vals.as_raw(), - T::DTYPE, + T::DTYPE.as_raw(), stream.as_raw(), ) }; @@ -118,8 +119,16 @@ impl Array { error::install(); let mut out = unsafe { sys::mlx_array_new() }; // SAFETY: stream is valid; the result is written into `out`. - let status = - unsafe { sys::mlx_arange(&mut out, start, stop, step, T::DTYPE, stream.as_raw()) }; + let status = unsafe { + sys::mlx_arange( + &mut out, + start, + stop, + step, + T::DTYPE.as_raw(), + stream.as_raw(), + ) + }; Self::from_op(out, status) } @@ -142,6 +151,16 @@ impl Array { (0..ndim).map(|i| unsafe { *ptr.add(i) }).collect() } + /// Element type of the array. + /// + /// MLX decides this itself: constructors take it from the Rust type, but ops + /// promote (an `i32` array times a `f32` one is `f32`) and comparisons always + /// give [`Dtype::Bool`], so this is the way to check what an op produced. + pub fn dtype(&self) -> Dtype { + // SAFETY: mlx_array_dtype reads a valid handle. + Dtype::from_raw(unsafe { sys::mlx_array_dtype(self.handle) }) + } + /// Strides of the array, in elements (not bytes), one per dimension. pub fn strides(&self) -> Vec { let ndim = self.ndim(); @@ -246,12 +265,11 @@ impl Array { /// Panics if `T::DTYPE` does not match the array's dtype. Assumes the array /// is already evaluated and row-contiguous. fn read_buffer(&self) -> Vec { - // SAFETY: mlx_array_dtype reads a valid handle. - let dtype = unsafe { sys::mlx_array_dtype(self.handle) }; + let dtype = self.dtype(); assert_eq!( dtype, T::DTYPE, - "array dtype does not match requested element type" + "array dtype {dtype} does not match requested element type" ); let len = self.size(); // `from_raw_parts` requires a non-null, aligned pointer even for a @@ -285,7 +303,8 @@ impl Array { error::install(); let mut out = unsafe { sys::mlx_array_new() }; // SAFETY: handle/stream are valid; the result is written into `out`. - let status = unsafe { sys::mlx_astype(&mut out, self.handle, T::DTYPE, stream.as_raw()) }; + let status = + unsafe { sys::mlx_astype(&mut out, self.handle, T::DTYPE.as_raw(), stream.as_raw()) }; Self::from_op(out, status) } @@ -379,6 +398,86 @@ impl Array { self.binary_op(other, stream, sys::mlx_minimum) } + /// Elementwise `self == other`, as a `bool` array. + /// + /// Like the arithmetic ops these broadcast, so comparing against a + /// [`from_scalar`](Self::from_scalar) array gives a mask over every element. + /// Read the result back with `to_vec::()`. + /// + /// These are named after the MLX operations rather than spelled `eq`/`lt`, + /// because they are elementwise and return an array — they are not the + /// whole-array `bool` answer that `PartialEq`/`PartialOrd` would imply. For + /// that, see [`array_equal`](Self::array_equal). + pub fn equal(&self, other: &Array, stream: &Stream) -> Result { + self.binary_op(other, stream, sys::mlx_equal) + } + + /// Elementwise `self != other`, as a `bool` array. + pub fn not_equal(&self, other: &Array, stream: &Stream) -> Result { + self.binary_op(other, stream, sys::mlx_not_equal) + } + + /// Elementwise `self > other`, as a `bool` array. + pub fn greater(&self, other: &Array, stream: &Stream) -> Result { + self.binary_op(other, stream, sys::mlx_greater) + } + + /// Elementwise `self >= other`, as a `bool` array. + pub fn greater_equal(&self, other: &Array, stream: &Stream) -> Result { + self.binary_op(other, stream, sys::mlx_greater_equal) + } + + /// Elementwise `self < other`, as a `bool` array. + pub fn less(&self, other: &Array, stream: &Stream) -> Result { + self.binary_op(other, stream, sys::mlx_less) + } + + /// Elementwise `self <= other`, as a `bool` array. + pub fn less_equal(&self, other: &Array, stream: &Stream) -> Result { + self.binary_op(other, stream, sys::mlx_less_equal) + } + + /// Elementwise logical and, as a `bool` array. + /// + /// Non-`bool` operands are compared against zero first, so this is a + /// truthiness test rather than a bitwise one. + pub fn logical_and(&self, other: &Array, stream: &Stream) -> Result { + self.binary_op(other, stream, sys::mlx_logical_and) + } + + /// Elementwise logical or, as a `bool` array. + pub fn logical_or(&self, other: &Array, stream: &Stream) -> Result { + self.binary_op(other, stream, sys::mlx_logical_or) + } + + /// Elementwise logical negation, as a `bool` array. + pub fn logical_not(&self, stream: &Stream) -> Result { + self.unary_op(stream, sys::mlx_logical_not) + } + + /// Whether the two arrays have the same shape and equal elements, as a + /// 0-dimensional `bool` array. + /// + /// This is the whole-array answer, in contrast to the elementwise + /// [`equal`](Self::equal). Arrays of different shapes are unequal — unlike + /// `equal`, nothing is broadcast. With `equal_nan == true` two NaNs in the + /// same position count as equal. + pub fn array_equal(&self, other: &Array, equal_nan: bool, stream: &Stream) -> Result { + error::install(); + let mut out = unsafe { sys::mlx_array_new() }; + // SAFETY: all handles are valid; the result is written into `out`. + let status = unsafe { + sys::mlx_array_equal( + &mut out, + self.handle, + other.as_raw(), + equal_nan, + stream.as_raw(), + ) + }; + Self::from_op(out, status) + } + /// Matrix multiplication: `self @ other`. /// /// Unlike the elementwise ops, this contracts the last axis of `self` @@ -415,6 +514,43 @@ impl Array { self.reduce_op(keepdims, stream, sys::mlx_mean) } + /// Maximum of all elements, returning a scalar array. + /// + /// With `keepdims == false` the result is 0-dimensional. + pub fn max(&self, keepdims: bool, stream: &Stream) -> Result { + self.reduce_op(keepdims, stream, sys::mlx_max) + } + + /// Minimum of all elements, returning a scalar array. + /// + /// With `keepdims == false` the result is 0-dimensional. + pub fn min(&self, keepdims: bool, stream: &Stream) -> Result { + self.reduce_op(keepdims, stream, sys::mlx_min) + } + + /// Product of all elements, returning a scalar array. + /// + /// With `keepdims == false` the result is 0-dimensional. + pub fn prod(&self, keepdims: bool, stream: &Stream) -> Result { + self.reduce_op(keepdims, stream, sys::mlx_prod) + } + + /// Whether every element is true (nonzero), as a `bool` array. + /// + /// With `keepdims == false` the result is 0-dimensional. Empty arrays reduce + /// to `true`, the identity of logical and. + pub fn all(&self, keepdims: bool, stream: &Stream) -> Result { + self.reduce_op(keepdims, stream, sys::mlx_all) + } + + /// Whether any element is true (nonzero), as a `bool` array. + /// + /// With `keepdims == false` the result is 0-dimensional. Empty arrays reduce + /// to `false`, the identity of logical or. + pub fn any(&self, keepdims: bool, stream: &Stream) -> Result { + self.reduce_op(keepdims, stream, sys::mlx_any) + } + /// Sum over the given axes. /// /// With `keepdims == false` the reduced axes are removed; otherwise they @@ -455,6 +591,22 @@ impl Array { self.reduce_axes_op(axes, keepdims, stream, sys::mlx_prod_axes) } + /// Whether every element is true (nonzero) over the given axes. + /// + /// With `keepdims == false` the reduced axes are removed; otherwise they + /// are kept with size 1. + pub fn all_axes(&self, axes: &[i32], keepdims: bool, stream: &Stream) -> Result { + self.reduce_axes_op(axes, keepdims, stream, sys::mlx_all_axes) + } + + /// Whether any element is true (nonzero) over the given axes. + /// + /// With `keepdims == false` the reduced axes are removed; otherwise they + /// are kept with size 1. + pub fn any_axes(&self, axes: &[i32], keepdims: bool, stream: &Stream) -> Result { + self.reduce_axes_op(axes, keepdims, stream, sys::mlx_any_axes) + } + /// Returns a new array with the same data reinterpreted as `shape`. /// /// The product of `shape` must equal [`size`](Self::size). @@ -490,7 +642,7 @@ impl Array { /// constructors. fn fill_op( shape: &[i32], - dtype: sys::mlx_dtype, + dtype: Dtype, stream: &Stream, op: unsafe extern "C" fn( *mut sys::mlx_array, @@ -505,7 +657,15 @@ impl Array { let mut out = unsafe { sys::mlx_array_new() }; // SAFETY: `shape_ptr`/`shape.len()` describe a valid slice (or null/0); // stream is valid; `op` writes the result into `out`. - let status = unsafe { op(&mut out, shape_ptr, shape.len(), dtype, stream.as_raw()) }; + let status = unsafe { + op( + &mut out, + shape_ptr, + shape.len(), + dtype.as_raw(), + stream.as_raw(), + ) + }; Self::from_op(out, status) } @@ -1095,6 +1255,226 @@ mod tests { ); } + #[test] + fn full_reductions_over_all_elements() { + let s = Stream::cpu(); + // Reduces across every axis, not just the last one. + let a = Array::from_slice(&[3.0f32, 1.0, 4.0, 1.0, 5.0, 9.0], &[2, 3]); + assert_eq!(a.max(false, &s).unwrap().item::(), 9.0); + assert_eq!(a.min(false, &s).unwrap().item::(), 1.0); + + let b = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2]); + assert_eq!(b.prod(false, &s).unwrap().item::(), 24.0); + } + + #[test] + fn full_reductions_respect_keepdims() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2]); + + // Without keepdims the rank collapses to 0... + assert_eq!(a.max(false, &s).unwrap().ndim(), 0); + // ...with it, every reduced axis is kept at size 1. + assert_eq!(a.max(true, &s).unwrap().shape(), vec![1, 1]); + assert_eq!(a.min(true, &s).unwrap().shape(), vec![1, 1]); + assert_eq!(a.prod(true, &s).unwrap().shape(), vec![1, 1]); + } + + #[test] + fn full_reductions_agree_with_axis_versions() { + let s = Stream::cpu(); + let a = Array::from_slice(&[3.0f32, 1.0, 4.0, 1.0, 5.0, 9.0], &[2, 3]); + // Reducing over all axes explicitly must match the full reduction. + assert_eq!( + a.max(false, &s).unwrap().item::(), + a.max_axes(&[0, 1], false, &s).unwrap().item::() + ); + assert_eq!( + a.min(false, &s).unwrap().item::(), + a.min_axes(&[0, 1], false, &s).unwrap().item::() + ); + assert_eq!( + a.prod(false, &s).unwrap().item::(), + a.prod_axes(&[0, 1], false, &s).unwrap().item::() + ); + } + + #[test] + fn dtype_reports_the_element_type() { + let s = Stream::cpu(); + assert_eq!( + Array::from_slice(&[1.0f32, 2.0], &[2]).dtype(), + Dtype::Float32 + ); + assert_eq!(Array::from_slice(&[1i32, 2], &[2]).dtype(), Dtype::Int32); + assert_eq!(Array::from_scalar(true).dtype(), Dtype::Bool); + assert_eq!(Array::zeros::(&[2], &s).unwrap().dtype(), Dtype::Uint8); + + // Ops decide the dtype themselves: astype converts, and mixing widens. + let ints = Array::from_slice(&[1i32, 2], &[2]); + assert_eq!(ints.astype::(&s).unwrap().dtype(), Dtype::Float32); + let floats = Array::from_slice(&[0.5f32, 0.5], &[2]); + assert_eq!(ints.add(&floats, &s).unwrap().dtype(), Dtype::Float32); + } + + #[test] + #[should_panic(expected = "array dtype float32 does not match")] + fn reading_the_wrong_element_type_panics() { + let _ = Array::from_slice(&[1.0f32, 2.0], &[2]).to_vec::(); + } + + #[test] + fn comparisons_produce_bool_masks() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0, 3.0], &[3]); + let b = Array::from_slice(&[3.0f32, 2.0, 1.0], &[3]); + + let eq = a.equal(&b, &s).unwrap(); + // Comparing floats gives bool, not float. + assert_eq!(eq.dtype(), Dtype::Bool); + assert_eq!(eq.to_vec::(), vec![false, true, false]); + + assert_eq!( + a.not_equal(&b, &s).unwrap().to_vec::(), + vec![true, false, true] + ); + assert_eq!( + a.greater(&b, &s).unwrap().to_vec::(), + vec![false, false, true] + ); + assert_eq!( + a.greater_equal(&b, &s).unwrap().to_vec::(), + vec![false, true, true] + ); + assert_eq!( + a.less(&b, &s).unwrap().to_vec::(), + vec![true, false, false] + ); + assert_eq!( + a.less_equal(&b, &s).unwrap().to_vec::(), + vec![true, true, false] + ); + } + + #[test] + fn comparisons_broadcast_against_a_scalar() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2]); + let mask = a.greater(&Array::from_scalar(2.0f32), &s).unwrap(); + // The scalar stretches over both axes, so the mask keeps `a`'s shape. + assert_eq!(mask.shape(), vec![2, 2]); + assert_eq!(mask.to_vec::(), vec![false, false, true, true]); + } + + #[test] + fn logical_ops_combine_masks() { + let s = Stream::cpu(); + let x = Array::from_slice(&[true, true, false, false], &[4]); + let y = Array::from_slice(&[true, false, true, false], &[4]); + + assert_eq!( + x.logical_and(&y, &s).unwrap().to_vec::(), + vec![true, false, false, false] + ); + assert_eq!( + x.logical_or(&y, &s).unwrap().to_vec::(), + vec![true, true, true, false] + ); + assert_eq!( + x.logical_not(&s).unwrap().to_vec::(), + vec![false, false, true, true] + ); + } + + #[test] + fn logical_ops_test_truthiness_of_numbers() { + let s = Stream::cpu(); + // 2.0 is neither 1 nor 0: a bitwise `and` would give 0, truthiness gives true. + let a = Array::from_slice(&[2.0f32, 0.0], &[2]); + let b = Array::from_slice(&[4.0f32, 7.0], &[2]); + assert_eq!( + a.logical_and(&b, &s).unwrap().to_vec::(), + vec![true, false] + ); + assert_eq!( + a.logical_not(&s).unwrap().to_vec::(), + vec![false, true] + ); + } + + #[test] + fn all_and_any_reduce_masks_to_one_answer() { + let s = Stream::cpu(); + let mixed = Array::from_slice(&[true, false], &[2]); + assert!(!mixed.all(false, &s).unwrap().item::()); + assert!(mixed.any(false, &s).unwrap().item::()); + + let all_true = Array::from_slice(&[true, true], &[2]); + assert!(all_true.all(false, &s).unwrap().item::()); + + let none = Array::from_slice(&[false, false], &[2]); + assert!(!none.any(false, &s).unwrap().item::()); + + // The reduction is over every axis, and `keepdims` behaves as elsewhere. + let grid = Array::from_slice(&[true, false, true, true], &[2, 2]); + assert!(!grid.all(false, &s).unwrap().item::()); + assert_eq!(grid.any(true, &s).unwrap().shape(), vec![1, 1]); + } + + #[test] + fn all_and_any_over_axes_reduce_per_row() { + let s = Stream::cpu(); + let grid = Array::from_slice(&[true, false, true, true], &[2, 2]); + // Row 0 is mixed, row 1 is all true. + assert_eq!( + grid.all_axes(&[1], false, &s).unwrap().to_vec::(), + vec![false, true] + ); + assert_eq!( + grid.any_axes(&[1], false, &s).unwrap().to_vec::(), + vec![true, true] + ); + // Column 0 is all true, column 1 is mixed. + assert_eq!( + grid.all_axes(&[0], false, &s).unwrap().to_vec::(), + vec![true, false] + ); + } + + #[test] + fn empty_reductions_return_the_identity() { + let s = Stream::cpu(); + let empty = Array::from_slice::(&[], &[0]); + assert!(empty.all(false, &s).unwrap().item::()); + assert!(!empty.any(false, &s).unwrap().item::()); + } + + #[test] + fn array_equal_compares_whole_arrays() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, 2.0, 3.0], &[3]); + let same = Array::from_slice(&[1.0f32, 2.0, 3.0], &[3]); + let different = Array::from_slice(&[1.0f32, 9.0, 3.0], &[3]); + + assert!(a.array_equal(&same, false, &s).unwrap().item::()); + assert!(!a.array_equal(&different, false, &s).unwrap().item::()); + // Unlike `equal`, nothing is broadcast: a different shape is unequal + // even when the elements line up. + let row = Array::from_slice(&[1.0f32, 2.0, 3.0], &[1, 3]); + assert!(!a.array_equal(&row, false, &s).unwrap().item::()); + assert_eq!(a.equal(&row, &s).unwrap().shape(), vec![1, 3]); + } + + #[test] + fn array_equal_can_treat_nans_as_equal() { + let s = Stream::cpu(); + let a = Array::from_slice(&[1.0f32, f32::NAN], &[2]); + let b = Array::from_slice(&[1.0f32, f32::NAN], &[2]); + // NaN != NaN, so the default comparison fails. + assert!(!a.array_equal(&b, false, &s).unwrap().item::()); + assert!(a.array_equal(&b, true, &s).unwrap().item::()); + } + #[test] #[should_panic(expected = "does not match shape product")] fn mismatched_len_and_shape_panics() { diff --git a/crates/mlxcore/src/dtype.rs b/crates/mlxcore/src/dtype.rs index 55851e3..1630d6a 100644 --- a/crates/mlxcore/src/dtype.rs +++ b/crates/mlxcore/src/dtype.rs @@ -1,19 +1,126 @@ //! Mapping between Rust primitive types and MLX data types. +use std::fmt; + use mlxcore_sys as sys; mod sealed { pub trait Sealed {} } -/// A Rust type that has a corresponding MLX [`mlx_dtype`](sys::mlx_dtype). +/// The element type of an [`Array`](crate::Array). +/// +/// Every MLX dtype has a variant here, including the three with no Rust +/// primitive to match them ([`Float16`](Dtype::Float16), +/// [`Bfloat16`](Dtype::Bfloat16), [`Complex64`](Dtype::Complex64)) — an array +/// can carry those even though we cannot yet read or construct them, so +/// [`Array::dtype`](crate::Array::dtype) must be able to name them. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Dtype { + Bool, + Uint8, + Uint16, + Uint32, + Uint64, + Int8, + Int16, + Int32, + Int64, + Float16, + Float32, + Float64, + Bfloat16, + Complex64, +} + +impl Dtype { + /// The raw mlx-c enum value, for passing across the FFI boundary. + pub(crate) fn as_raw(self) -> sys::mlx_dtype { + match self { + Dtype::Bool => sys::MLX_BOOL, + Dtype::Uint8 => sys::MLX_UINT8, + Dtype::Uint16 => sys::MLX_UINT16, + Dtype::Uint32 => sys::MLX_UINT32, + Dtype::Uint64 => sys::MLX_UINT64, + Dtype::Int8 => sys::MLX_INT8, + Dtype::Int16 => sys::MLX_INT16, + Dtype::Int32 => sys::MLX_INT32, + Dtype::Int64 => sys::MLX_INT64, + Dtype::Float16 => sys::MLX_FLOAT16, + Dtype::Float32 => sys::MLX_FLOAT32, + Dtype::Float64 => sys::MLX_FLOAT64, + Dtype::Bfloat16 => sys::MLX_BFLOAT16, + Dtype::Complex64 => sys::MLX_COMPLEX64, + } + } + + /// Converts a raw mlx-c enum value. + /// + /// # Panics + /// Panics on a value MLX did not define when these bindings were generated, + /// which would mean the linked MLX added a dtype we cannot name. + pub(crate) fn from_raw(raw: sys::mlx_dtype) -> Self { + match raw { + sys::MLX_BOOL => Dtype::Bool, + sys::MLX_UINT8 => Dtype::Uint8, + sys::MLX_UINT16 => Dtype::Uint16, + sys::MLX_UINT32 => Dtype::Uint32, + sys::MLX_UINT64 => Dtype::Uint64, + sys::MLX_INT8 => Dtype::Int8, + sys::MLX_INT16 => Dtype::Int16, + sys::MLX_INT32 => Dtype::Int32, + sys::MLX_INT64 => Dtype::Int64, + sys::MLX_FLOAT16 => Dtype::Float16, + sys::MLX_FLOAT32 => Dtype::Float32, + sys::MLX_FLOAT64 => Dtype::Float64, + sys::MLX_BFLOAT16 => Dtype::Bfloat16, + sys::MLX_COMPLEX64 => Dtype::Complex64, + other => panic!("unknown mlx dtype: {other}"), + } + } + + /// The size of one element in bytes. + pub fn size(self) -> usize { + match self { + Dtype::Bool | Dtype::Uint8 | Dtype::Int8 => 1, + Dtype::Uint16 | Dtype::Int16 | Dtype::Float16 | Dtype::Bfloat16 => 2, + Dtype::Uint32 | Dtype::Int32 | Dtype::Float32 => 4, + Dtype::Uint64 | Dtype::Int64 | Dtype::Float64 | Dtype::Complex64 => 8, + } + } +} + +/// Formats as MLX's own name for the dtype, e.g. `float32`, not `Float32`. +impl fmt::Display for Dtype { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let name = match self { + Dtype::Bool => "bool", + Dtype::Uint8 => "uint8", + Dtype::Uint16 => "uint16", + Dtype::Uint32 => "uint32", + Dtype::Uint64 => "uint64", + Dtype::Int8 => "int8", + Dtype::Int16 => "int16", + Dtype::Int32 => "int32", + Dtype::Int64 => "int64", + Dtype::Float16 => "float16", + Dtype::Float32 => "float32", + Dtype::Float64 => "float64", + Dtype::Bfloat16 => "bfloat16", + Dtype::Complex64 => "complex64", + }; + f.write_str(name) + } +} + +/// A Rust type that has a corresponding MLX [`Dtype`]. /// /// This trait is sealed: it can only be implemented for the primitive types /// MLX supports, so `T::DTYPE` is always a valid dtype and the accessors below /// always match it. pub trait ArrayElement: sealed::Sealed + Copy + Default { /// The MLX dtype corresponding to this Rust type. - const DTYPE: sys::mlx_dtype; + const DTYPE: Dtype; /// Reads the value of a scalar (single-element) array as this type. /// @@ -34,7 +141,7 @@ macro_rules! impl_array_element { $( impl sealed::Sealed for $rust {} impl ArrayElement for $rust { - const DTYPE: sys::mlx_dtype = $dtype; + const DTYPE: Dtype = $dtype; unsafe fn read_item(arr: sys::mlx_array) -> Self { let mut out = <$rust>::default(); @@ -56,17 +163,17 @@ macro_rules! impl_array_element { // without a stable Rust equivalent (float16, bfloat16, complex64) are left for // dedicated newtypes later. impl_array_element! { - bool => sys::MLX_BOOL, sys::mlx_array_item_bool, sys::mlx_array_data_bool, - u8 => sys::MLX_UINT8, sys::mlx_array_item_uint8, sys::mlx_array_data_uint8, - u16 => sys::MLX_UINT16, sys::mlx_array_item_uint16, sys::mlx_array_data_uint16, - u32 => sys::MLX_UINT32, sys::mlx_array_item_uint32, sys::mlx_array_data_uint32, - u64 => sys::MLX_UINT64, sys::mlx_array_item_uint64, sys::mlx_array_data_uint64, - i8 => sys::MLX_INT8, sys::mlx_array_item_int8, sys::mlx_array_data_int8, - i16 => sys::MLX_INT16, sys::mlx_array_item_int16, sys::mlx_array_data_int16, - i32 => sys::MLX_INT32, sys::mlx_array_item_int32, sys::mlx_array_data_int32, - i64 => sys::MLX_INT64, sys::mlx_array_item_int64, sys::mlx_array_data_int64, - f32 => sys::MLX_FLOAT32, sys::mlx_array_item_float32, sys::mlx_array_data_float32, - f64 => sys::MLX_FLOAT64, sys::mlx_array_item_float64, sys::mlx_array_data_float64, + bool => Dtype::Bool, sys::mlx_array_item_bool, sys::mlx_array_data_bool, + u8 => Dtype::Uint8, sys::mlx_array_item_uint8, sys::mlx_array_data_uint8, + u16 => Dtype::Uint16, sys::mlx_array_item_uint16, sys::mlx_array_data_uint16, + u32 => Dtype::Uint32, sys::mlx_array_item_uint32, sys::mlx_array_data_uint32, + u64 => Dtype::Uint64, sys::mlx_array_item_uint64, sys::mlx_array_data_uint64, + i8 => Dtype::Int8, sys::mlx_array_item_int8, sys::mlx_array_data_int8, + i16 => Dtype::Int16, sys::mlx_array_item_int16, sys::mlx_array_data_int16, + i32 => Dtype::Int32, sys::mlx_array_item_int32, sys::mlx_array_data_int32, + i64 => Dtype::Int64, sys::mlx_array_item_int64, sys::mlx_array_data_int64, + f32 => Dtype::Float32, sys::mlx_array_item_float32, sys::mlx_array_data_float32, + f64 => Dtype::Float64, sys::mlx_array_item_float64, sys::mlx_array_data_float64, } #[cfg(test)] @@ -75,10 +182,55 @@ mod tests { #[test] fn maps_rust_types_to_expected_dtypes() { - assert_eq!(::DTYPE, sys::MLX_FLOAT32); - assert_eq!(::DTYPE, sys::MLX_FLOAT64); - assert_eq!(::DTYPE, sys::MLX_INT32); - assert_eq!(::DTYPE, sys::MLX_UINT8); - assert_eq!(::DTYPE, sys::MLX_BOOL); + assert_eq!(::DTYPE, Dtype::Float32); + assert_eq!(::DTYPE, Dtype::Float64); + assert_eq!(::DTYPE, Dtype::Int32); + assert_eq!(::DTYPE, Dtype::Uint8); + assert_eq!(::DTYPE, Dtype::Bool); + } + + #[test] + fn dtypes_round_trip_through_the_raw_enum() { + const ALL: [Dtype; 14] = [ + Dtype::Bool, + Dtype::Uint8, + Dtype::Uint16, + Dtype::Uint32, + Dtype::Uint64, + Dtype::Int8, + Dtype::Int16, + Dtype::Int32, + Dtype::Int64, + Dtype::Float16, + Dtype::Float32, + Dtype::Float64, + Dtype::Bfloat16, + Dtype::Complex64, + ]; + for dtype in ALL { + assert_eq!(Dtype::from_raw(dtype.as_raw()), dtype); + } + } + + #[test] + fn raw_values_match_the_c_enum() { + assert_eq!(Dtype::Float32.as_raw(), sys::MLX_FLOAT32); + assert_eq!(Dtype::Bfloat16.as_raw(), sys::MLX_BFLOAT16); + assert_eq!(Dtype::from_raw(sys::MLX_COMPLEX64), Dtype::Complex64); + } + + #[test] + #[should_panic(expected = "unknown mlx dtype")] + fn unknown_raw_dtype_panics() { + let _ = Dtype::from_raw(sys::MLX_COMPLEX64 + 1); + } + + #[test] + fn displays_mlx_names_and_element_sizes() { + assert_eq!(Dtype::Float32.to_string(), "float32"); + assert_eq!(Dtype::Bfloat16.to_string(), "bfloat16"); + assert_eq!(Dtype::Bool.size(), 1); + assert_eq!(Dtype::Float32.size(), 4); + assert_eq!(Dtype::Complex64.size(), 8); } } diff --git a/crates/mlxcore/src/lib.rs b/crates/mlxcore/src/lib.rs index aeb9517..a0d0713 100644 --- a/crates/mlxcore/src/lib.rs +++ b/crates/mlxcore/src/lib.rs @@ -22,7 +22,7 @@ mod stream; pub mod random; pub use array::Array; -pub use dtype::ArrayElement; +pub use dtype::{ArrayElement, Dtype}; pub use error::{Error, Result}; pub use stream::Stream; diff --git a/crates/mlxcore/src/random.rs b/crates/mlxcore/src/random.rs index 7e5075b..53d3266 100644 --- a/crates/mlxcore/src/random.rs +++ b/crates/mlxcore/src/random.rs @@ -109,7 +109,7 @@ pub fn normal( &mut out, as_ffi_ptr(shape), shape.len(), - T::DTYPE, + T::DTYPE.as_raw(), loc, scale, Key::handle(key), @@ -144,7 +144,7 @@ pub fn uniform( high.as_raw(), as_ffi_ptr(shape), shape.len(), - T::DTYPE, + T::DTYPE.as_raw(), Key::handle(key), stream.as_raw(), )