mlx_c_rust/types/
array.rs

1use crate::{
2    MlxCError,
3    error::{check_status, clear_captured_mlx_error, take_captured_mlx_error},
4    raw,
5};
6
7/// MLX array element types exposed without leaking generated C declarations.
8#[derive(Clone, Copy, Debug, Eq, PartialEq)]
9pub enum MlxDtype {
10    Bool,
11    UInt8,
12    UInt16,
13    UInt32,
14    UInt64,
15    Int8,
16    Int16,
17    Int32,
18    Int64,
19    Float16,
20    Float32,
21    Float64,
22    BFloat16,
23    Complex64,
24}
25
26impl MlxDtype {
27    fn from_raw(raw_dtype: raw::mlx_dtype) -> Self {
28        match raw_dtype {
29            raw::mlx_dtype__MLX_BOOL => Self::Bool,
30            raw::mlx_dtype__MLX_UINT8 => Self::UInt8,
31            raw::mlx_dtype__MLX_UINT16 => Self::UInt16,
32            raw::mlx_dtype__MLX_UINT32 => Self::UInt32,
33            raw::mlx_dtype__MLX_UINT64 => Self::UInt64,
34            raw::mlx_dtype__MLX_INT8 => Self::Int8,
35            raw::mlx_dtype__MLX_INT16 => Self::Int16,
36            raw::mlx_dtype__MLX_INT32 => Self::Int32,
37            raw::mlx_dtype__MLX_INT64 => Self::Int64,
38            raw::mlx_dtype__MLX_FLOAT16 => Self::Float16,
39            raw::mlx_dtype__MLX_FLOAT32 => Self::Float32,
40            raw::mlx_dtype__MLX_FLOAT64 => Self::Float64,
41            raw::mlx_dtype__MLX_BFLOAT16 => Self::BFloat16,
42            raw::mlx_dtype__MLX_COMPLEX64 => Self::Complex64,
43            _ => Self::Bool,
44        }
45    }
46
47    /// The raw MLX dtype tag for raw-operation builders outside this crate.
48    pub const fn to_raw(self) -> raw::mlx_dtype {
49        match self {
50            Self::Bool => raw::mlx_dtype__MLX_BOOL,
51            Self::UInt8 => raw::mlx_dtype__MLX_UINT8,
52            Self::UInt16 => raw::mlx_dtype__MLX_UINT16,
53            Self::UInt32 => raw::mlx_dtype__MLX_UINT32,
54            Self::UInt64 => raw::mlx_dtype__MLX_UINT64,
55            Self::Int8 => raw::mlx_dtype__MLX_INT8,
56            Self::Int16 => raw::mlx_dtype__MLX_INT16,
57            Self::Int32 => raw::mlx_dtype__MLX_INT32,
58            Self::Int64 => raw::mlx_dtype__MLX_INT64,
59            Self::Float16 => raw::mlx_dtype__MLX_FLOAT16,
60            Self::Float32 => raw::mlx_dtype__MLX_FLOAT32,
61            Self::Float64 => raw::mlx_dtype__MLX_FLOAT64,
62            Self::BFloat16 => raw::mlx_dtype__MLX_BFLOAT16,
63            Self::Complex64 => raw::mlx_dtype__MLX_COMPLEX64,
64        }
65    }
66}
67
68/// Owned MLX array handle released exactly once through the official C API.
69#[derive(Debug)]
70pub struct MlxArray {
71    pub(crate) raw_array: raw::mlx_array,
72}
73
74impl MlxArray {
75    /// Creates a float32 array from host values, the way compiled-graph
76    /// builders embed scalar constants (the softplus graph's ln 2, the
77    /// verification window's attention-mask boundary values).
78    pub fn from_f32(values: &[f32], shape: &[i32]) -> Result<Self, MlxCError> {
79        validate_shape(values.len(), shape)?;
80        let dimension_count = i32::try_from(shape.len()).map_err(|_| MlxCError {
81            operation: "create an MLX float32 array",
82            description: "array rank exceeds the C API integer range".to_owned(),
83        })?;
84        // SAFETY: The slices remain valid for this copying constructor, their
85        // lengths were validated, and the returned handle enters RAII ownership.
86        let raw_array = unsafe {
87            raw::mlx_array_new_data(
88                values.as_ptr().cast(),
89                shape.as_ptr(),
90                dimension_count,
91                raw::mlx_dtype__MLX_FLOAT32,
92            )
93        };
94        let array = Self { raw_array };
95        array.require_populated("create an MLX float32 array")?;
96        Ok(array)
97    }
98
99    pub fn from_i32(values: &[i32], shape: &[i32]) -> Result<Self, MlxCError> {
100        validate_shape(values.len(), shape)?;
101        let dimension_count = i32::try_from(shape.len()).map_err(|_| MlxCError {
102            operation: "create an MLX int32 array",
103            description: "array rank exceeds the C API integer range".to_owned(),
104        })?;
105        // SAFETY: The slices remain valid for this copying constructor, their
106        // lengths were validated, and the returned handle enters RAII ownership.
107        let raw_array = unsafe {
108            raw::mlx_array_new_data(
109                values.as_ptr().cast(),
110                shape.as_ptr(),
111                dimension_count,
112                raw::mlx_dtype__MLX_INT32,
113            )
114        };
115        let array = Self { raw_array };
116        array.require_populated("create an MLX int32 array")?;
117        Ok(array)
118    }
119
120    pub fn from_u32(values: &[u32], shape: &[i32]) -> Result<Self, MlxCError> {
121        validate_shape(values.len(), shape)?;
122        let dimension_count = i32::try_from(shape.len()).map_err(|_| MlxCError {
123            operation: "create an MLX uint32 array",
124            description: "array rank exceeds the C API integer range".to_owned(),
125        })?;
126        // SAFETY: The slices remain valid for this copying constructor, their
127        // lengths were validated, and the returned handle enters RAII ownership.
128        let raw_array = unsafe {
129            raw::mlx_array_new_data(
130                values.as_ptr().cast(),
131                shape.as_ptr(),
132                dimension_count,
133                raw::mlx_dtype__MLX_UINT32,
134            )
135        };
136        let array = Self { raw_array };
137        array.require_populated("create an MLX uint32 array")?;
138        Ok(array)
139    }
140
141    pub fn empty() -> Self {
142        // SAFETY: The runtime error handler is installed before model loading,
143        // and the official API's empty placeholder is immediately placed under
144        // RAII ownership before a fallible operation populates it.
145        let raw_array = unsafe { raw::mlx_array_new() };
146        Self { raw_array }
147    }
148
149    pub fn is_empty(&self) -> bool {
150        self.raw_array.ctx.is_null()
151    }
152
153    pub const fn raw(&self) -> raw::mlx_array {
154        self.raw_array
155    }
156
157    /// The official absent-optional-handle placeholder for raw MLX calls.
158    ///
159    /// Compiled-graph builders outside this crate mirror the runtime's own
160    /// operations and need the same empty convention for absent `freqs`,
161    /// `sinks`, and `global_scale` handles; the placeholder owns nothing and
162    /// is always valid to pass by value.
163    pub const fn empty_raw() -> raw::mlx_array {
164        raw::mlx_array {
165            ctx: std::ptr::null_mut(),
166        }
167    }
168
169    pub fn raw_mut(&mut self) -> *mut raw::mlx_array {
170        &mut self.raw_array
171    }
172
173    /// Validates that a returned handle is populated, translating a parked
174    /// payload-closure failure or the captured MLX-C description into a typed
175    /// failure.
176    ///
177    /// This is a pure translation: embedding-specific error classification
178    /// (for example an active memory ceiling) is embedding policy and happens
179    /// when the consumer converts this error at its boundary.
180    pub fn require_populated(&self, operation: &'static str) -> Result<(), MlxCError> {
181        if !self.is_empty() {
182            clear_captured_mlx_error();
183            crate::error::clear_closure_error();
184            return Ok(());
185        }
186        if let Some(failure) = crate::error::take_closure_error() {
187            clear_captured_mlx_error();
188            return Err(failure);
189        }
190        let description = take_captured_mlx_error()
191            .unwrap_or_else(|| "MLX returned an empty array handle".to_owned());
192        Err(MlxCError {
193            operation,
194            description,
195        })
196    }
197
198    /// Returns the dimensions of this live array.
199    #[must_use]
200    pub fn shape(&self) -> Vec<i32> {
201        // SAFETY: `self` exclusively owns a live MLX array handle.
202        let dimension_count = unsafe { raw::mlx_array_ndim(self.raw_array) };
203        if dimension_count == 0 {
204            return Vec::new();
205        }
206        // SAFETY: MLX keeps this shape storage alive with the array. The
207        // pointer addresses exactly `dimension_count` dimensions.
208        let shape_pointer = unsafe { raw::mlx_array_shape(self.raw_array) };
209        if shape_pointer.is_null() {
210            return Vec::new();
211        }
212        // SAFETY: The preceding MLX contract establishes the pointer length;
213        // this method copies the values before returning.
214        unsafe { std::slice::from_raw_parts(shape_pointer, dimension_count) }.to_vec()
215    }
216
217    /// Returns the array's element type.
218    #[must_use]
219    pub fn dtype(&self) -> MlxDtype {
220        // SAFETY: `self` exclusively owns a live MLX array handle.
221        MlxDtype::from_raw(unsafe { raw::mlx_array_dtype(self.raw_array) })
222    }
223
224    /// Returns the number of logical elements.
225    #[must_use]
226    pub fn element_count(&self) -> usize {
227        // SAFETY: `self` exclusively owns a live MLX array handle.
228        unsafe { raw::mlx_array_size(self.raw_array) }
229    }
230
231    /// Returns the logical tensor payload size in bytes.
232    #[must_use]
233    pub fn byte_count(&self) -> usize {
234        // SAFETY: `self` exclusively owns a live MLX array handle.
235        unsafe { raw::mlx_array_nbytes(self.raw_array) }
236    }
237
238    /// Materializes this lazy array and reports runtime failures without exit.
239    pub fn evaluate(&self) -> Result<(), MlxCError> {
240        // SAFETY: `self` owns a live handle and MLX evaluation does not retain
241        // any Rust borrow beyond this call.
242        let status = unsafe { raw::mlx_array_eval(self.raw_array) };
243        check_status(status, "evaluate an MLX array")
244    }
245
246    /// Evaluates and copies a contiguous float32 array into Rust-owned memory.
247    ///
248    /// MLX keeps a slice of an evaluated parent as a strided view over the
249    /// parent's buffer, and the raw data pointer ignores those strides, so
250    /// this readout is only valid for contiguous arrays. Read a strided view
251    /// through `MlxRuntime::array_to_vec_f32`, whose flat reshape forces one
252    /// contiguous materialization.
253    pub fn to_vec_f32(&self) -> Result<Vec<f32>, MlxCError> {
254        if self.dtype() != MlxDtype::Float32 {
255            return Err(MlxCError {
256                operation: "copy an MLX float32 array",
257                description: "array dtype is not float32".to_owned(),
258            });
259        }
260        self.evaluate()?;
261        let element_count = self.element_count();
262        if element_count == 0 {
263            return Ok(Vec::new());
264        }
265        // SAFETY: Evaluation materializes contiguous readable storage owned by
266        // the live array for at least `element_count` float32 values.
267        let values_pointer = unsafe { raw::mlx_array_data_float32(self.raw_array) };
268        if values_pointer.is_null() {
269            return Err(MlxCError {
270                operation: "copy an MLX float32 array",
271                description: "MLX returned a null data pointer after evaluation".to_owned(),
272            });
273        }
274        // SAFETY: The evaluated MLX array establishes the pointer's exact
275        // element count; values are copied before the borrow ends.
276        Ok(unsafe { std::slice::from_raw_parts(values_pointer, element_count) }.to_vec())
277    }
278
279    pub fn to_vec_u32(&self) -> Result<Vec<u32>, MlxCError> {
280        if self.dtype() != MlxDtype::UInt32 {
281            return Err(MlxCError {
282                operation: "copy an MLX uint32 array",
283                description: "array dtype is not uint32".to_owned(),
284            });
285        }
286        self.evaluate()?;
287        self.copy_evaluated_u32_values()
288    }
289
290    /// Evaluates and copies a uint8 array into Rust-owned memory.
291    pub fn to_vec_u8(&self) -> Result<Vec<u8>, MlxCError> {
292        if self.dtype() != MlxDtype::UInt8 {
293            return Err(MlxCError {
294                operation: "copy an MLX uint8 array",
295                description: "array dtype is not uint8".to_owned(),
296            });
297        }
298        self.evaluate()?;
299        self.copy_evaluated_u8_values()
300    }
301
302    /// Copies an already evaluated contiguous uint8 array in one host transfer.
303    pub fn copy_evaluated_u8_values(&self) -> Result<Vec<u8>, MlxCError> {
304        if self.dtype() != MlxDtype::UInt8 {
305            return Err(MlxCError {
306                operation: "copy an evaluated MLX uint8 array",
307                description: "array dtype is not uint8".to_owned(),
308            });
309        }
310        let element_count = self.element_count();
311        if element_count == 0 {
312            return Ok(Vec::new());
313        }
314        // SAFETY: The caller evaluated contiguous storage that remains owned by
315        // this live array for at least `element_count` uint8 values.
316        let values_pointer = unsafe { raw::mlx_array_data_uint8(self.raw_array) };
317        if values_pointer.is_null() {
318            return Err(MlxCError {
319                operation: "copy an evaluated MLX uint8 array",
320                description: "MLX returned a null data pointer for an evaluated array".to_owned(),
321            });
322        }
323        // SAFETY: The evaluated array establishes the pointer's exact element
324        // count and the values are copied before the array can be released.
325        Ok(unsafe { std::slice::from_raw_parts(values_pointer, element_count) }.to_vec())
326    }
327
328    /// Copies an already evaluated uint32 array into Rust-owned memory.
329    ///
330    /// Callers must evaluate the array first. This split exists so performance
331    /// attribution can distinguish the synchronization wait from the small host copy.
332    pub fn copy_evaluated_u32_values(&self) -> Result<Vec<u32>, MlxCError> {
333        if self.dtype() != MlxDtype::UInt32 {
334            return Err(MlxCError {
335                operation: "copy an evaluated MLX uint32 array",
336                description: "array dtype is not uint32".to_owned(),
337            });
338        }
339        let element_count = self.element_count();
340        if element_count == 0 {
341            return Ok(Vec::new());
342        }
343        // SAFETY: Evaluation materializes contiguous readable storage owned by
344        // the live array for at least `element_count` uint32 values.
345        let values_pointer = unsafe { raw::mlx_array_data_uint32(self.raw_array) };
346        if values_pointer.is_null() {
347            return Err(MlxCError {
348                operation: "copy an evaluated MLX uint32 array",
349                description: "MLX returned a null data pointer for an evaluated array".to_owned(),
350            });
351        }
352        // SAFETY: The evaluated MLX array establishes the pointer's exact
353        // element count; values are copied before the borrow ends.
354        Ok(unsafe { std::slice::from_raw_parts(values_pointer, element_count) }.to_vec())
355    }
356
357    /// Evaluates a scalar array and copies its value as an unsigned token ID.
358    pub fn item_u32(&self) -> Result<u32, MlxCError> {
359        self.evaluate()?;
360        let mut scalar_value = 0_u32;
361        // SAFETY: The output pointer is valid writable storage and `self` owns
362        // a live evaluated array. MLX validates scalar shape and conversion.
363        let status = unsafe { raw::mlx_array_item_uint32(&mut scalar_value, self.raw_array) };
364        check_status(status, "read an MLX uint32 scalar")?;
365        Ok(scalar_value)
366    }
367
368    /// Creates a shared reference to the same underlying MLX data.
369    ///
370    /// MLX uses reference counting, so this is cheap and does not copy the
371    /// tensor data. The returned array shares the same lazy graph and
372    /// evaluated buffers as the source.
373    pub fn retain(&self) -> Result<Self, MlxCError> {
374        let mut retained_array = Self::empty();
375        // SAFETY: `retained_array` owns a live (empty) handle and `self` owns a
376        // live source handle. `mlx_array_set` copies the source reference into
377        // the destination.
378        let status = unsafe { raw::mlx_array_set(retained_array.raw_mut(), self.raw_array) };
379        check_status(status, "retain an MLX array")?;
380        if retained_array.is_empty() {
381            return Err(MlxCError {
382                operation: "retain an MLX array",
383                description: "MLX returned an empty handle after set".to_owned(),
384            });
385        }
386        Ok(retained_array)
387    }
388}
389
390fn validate_shape(element_count: usize, shape: &[i32]) -> Result<(), MlxCError> {
391    let shaped_element_count = shape.iter().try_fold(1_usize, |product, dimension| {
392        let dimension = usize::try_from(*dimension).ok()?;
393        product.checked_mul(dimension)
394    });
395    if shaped_element_count != Some(element_count) {
396        return Err(MlxCError {
397            operation: "create an MLX array",
398            description: "shape element count does not match the provided values".to_owned(),
399        });
400    }
401    Ok(())
402}
403
404impl Drop for MlxArray {
405    fn drop(&mut self) {
406        // SAFETY: This owner releases its live array exactly once and never
407        // accesses the handle afterward.
408        unsafe {
409            raw::mlx_array_free(self.raw_array);
410        }
411    }
412}
mlx_c_rust/types/array.rs

mlx_c_rust/types/
array.rs
