mlx_rs/
error.rs

1//! Custom error types and handler for the c ffi
2
3use crate::Dtype;
4use libc::strdup;
5use std::any::Any;
6use std::convert::Infallible;
7use std::ffi::NulError;
8use std::panic::Location;
9use std::sync::Once;
10use std::{cell::Cell, ffi::c_char};
11use thiserror::Error;
12
13/// Type alias for a `Result` with an `Exception` error type.
14pub type Result<T> = std::result::Result<T, Exception>;
15
16/// Error with io operations
17#[derive(Error, PartialEq, Debug)]
18pub enum IoError {
19    /// Path must point to a local file
20    #[error("Path must point to a local file")]
21    NotFile,
22
23    /// Path contains invalid UTF-8
24    #[error("Path contains invalid UTF-8")]
25    InvalidUtf8,
26
27    /// Path contains null bytes
28    #[error("Path contains null bytes")]
29    NullBytes,
30
31    /// No file extension found
32    #[error("No file extension found")]
33    NoExtension,
34
35    /// Unsupported file format
36    #[error("Unsupported file format")]
37    UnsupportedFormat,
38
39    /// Unable to open file
40    #[error("Unable to open file")]
41    UnableToOpenFile,
42
43    /// Unable to allocate memory
44    #[error("Unable to allocate memory")]
45    AllocationError,
46
47    /// Null error
48    #[error(transparent)]
49    NulError(#[from] NulError),
50
51    /// Error with unfalttening the loaded optimizer state
52    #[error(transparent)]
53    Unflatten(#[from] UnflattenError),
54
55    /// Error projecting optimizer state for serialization.
56    #[error(transparent)]
57    StateProjection(#[from] StateProjectionError),
58
59    /// Exception
60    #[error(transparent)]
61    Exception(#[from] Exception),
62}
63
64impl From<Infallible> for IoError {
65    fn from(_: Infallible) -> Self {
66        unreachable!()
67    }
68}
69
70impl From<RawException> for IoError {
71    #[track_caller]
72    fn from(e: RawException) -> Self {
73        let exception = Exception {
74            what: e.what,
75            location: Location::caller(),
76        };
77        Self::Exception(exception)
78    }
79}
80
81/// Error associated with `Array::try_as_slice()`
82#[derive(Debug, PartialEq, Error)]
83pub enum AsSliceError {
84    /// The array data is not contiguous in row-major order.
85    #[error(
86        "array data is not contiguous row-major; call `Array::contiguous()` before borrowing it as a slice"
87    )]
88    NotContiguous,
89
90    /// The underlying data pointer is null.
91    ///
92    /// This is likely because the array has not been evaluated yet.
93    #[error("The data pointer is null.")]
94    Null,
95
96    /// The output dtype does not match the data type of the array.
97    #[error("dtype mismatch: expected {expecting:?}, found {found:?}")]
98    DtypeMismatch {
99        /// The expected data type.
100        expecting: Dtype,
101
102        /// The actual data type
103        found: Dtype,
104    },
105
106    /// Exception
107    #[error(transparent)]
108    Exception(#[from] Exception),
109}
110
111/// Error with unflattening a loaded optimizer state
112#[derive(Debug, PartialEq, Error)]
113pub enum UnflattenError {
114    /// Expecting next (key, value) pair, found none
115    #[error("Expecting next (key, value) pair, found none")]
116    ExpectingNextPair,
117
118    /// The key is not in a valid format
119    #[error("Invalid key")]
120    InvalidKey,
121}
122
123/// Error with loading an optimizer state
124#[derive(Debug, PartialEq, Error)]
125pub enum OptimizerStateLoadError {
126    /// Error with io operations
127    #[error(transparent)]
128    Io(#[from] IoError),
129
130    /// Error with unflattening the optimizer state
131    #[error(transparent)]
132    Unflatten(#[from] UnflattenError),
133}
134
135/// Error declaring, restoring, or evaluating a keyed state projection.
136#[derive(Debug, PartialEq, Error)]
137pub enum StateProjectionError {
138    /// A projection declared the same key more than once.
139    #[error("duplicate state key {0}")]
140    DuplicateKey(String),
141
142    /// A projected key was absent from the supplied snapshot.
143    #[error("state snapshot is missing key {0}")]
144    MissingKey(String),
145
146    /// The supplied snapshot contains keys absent from the projection.
147    #[error("state snapshot contains unknown keys {0:?}")]
148    UnknownKeys(Vec<String>),
149
150    /// A required slot was absent in the supplied snapshot.
151    #[error("required state key {0} is absent")]
152    RequiredSlotAbsent(String),
153
154    /// A compiled state vector has the wrong number of present values.
155    #[error("state value count mismatch: expected {expected}, found {actual}")]
156    Cardinality {
157        /// The number of present values declared by the layout.
158        expected: usize,
159        /// The supplied value count.
160        actual: usize,
161    },
162
163    /// The upstream runtime rejected an operation needed to restore state.
164    #[error(transparent)]
165    Exception(#[from] Exception),
166}
167
168impl From<StateProjectionError> for Exception {
169    #[track_caller]
170    fn from(error: StateProjectionError) -> Self {
171        match error {
172            StateProjectionError::Exception(error) => error,
173            error => Self::custom(error.to_string()),
174        }
175    }
176}
177
178impl From<Infallible> for OptimizerStateLoadError {
179    fn from(_: Infallible) -> Self {
180        unreachable!()
181    }
182}
183
184/// Error converting an array or serialized tensor representation.
185#[derive(Debug, Error)]
186pub enum ConversionError {
187    /// The requested element type does not match the array dtype.
188    #[error("dtype mismatch: expected {expected:?}, found {actual:?}")]
189    DtypeMismatch {
190        /// The requested element type.
191        expected: Dtype,
192        /// The array element type.
193        actual: Dtype,
194    },
195
196    /// Scalar extraction was requested from a non-scalar array.
197    #[error("scalar extraction requires one element, found {actual}")]
198    NotScalar {
199        /// The array element count.
200        actual: usize,
201    },
202
203    /// The array cannot be borrowed as a contiguous slice.
204    #[error(transparent)]
205    ArraySlice(#[from] AsSliceError),
206
207    /// The upstream runtime rejected the conversion.
208    #[error(transparent)]
209    Exception(#[from] Exception),
210
211    /// The safetensors data type is not supported.
212    #[cfg(feature = "safetensors")]
213    #[error("The safetensors data type {0:?} is not supported.")]
214    SafeTensorDtype(safetensors::tensor::Dtype),
215
216    /// The MLX data type is not supported by safetensors.
217    #[cfg(feature = "safetensors")]
218    #[error("The mlx data type {0:?} is not supported.")]
219    MlxDtype(crate::Dtype),
220
221    /// Error casting the data buffer to `&[u8]`.
222    #[cfg(feature = "safetensors")]
223    #[error(transparent)]
224    PodCastError(#[from] bytemuck::PodCastError),
225
226    /// Error creating a safetensors tensor view.
227    #[cfg(feature = "safetensors")]
228    #[error(transparent)]
229    SafeTensorError(#[from] safetensors::tensor::SafeTensorError),
230}
231
232pub(crate) struct RawException {
233    pub(crate) what: String,
234}
235
236/// Exception. Most will come from the C API.
237#[derive(Debug, PartialEq, Error)]
238#[error("{what:?} at {location}")]
239pub struct Exception {
240    pub(crate) what: String,
241    pub(crate) location: &'static Location<'static>,
242}
243
244impl Exception {
245    /// The error message.
246    pub fn what(&self) -> &str {
247        &self.what
248    }
249
250    /// The location of the error.
251    ///
252    /// The location is obtained from `std::panic::Location::caller()` and points
253    /// to the location in the code where the error was created and not where it was
254    /// propagated.
255    pub fn location(&self) -> &'static Location<'static> {
256        self.location
257    }
258
259    /// Creates a new exception with the given message.
260    #[track_caller]
261    pub fn custom(what: impl Into<String>) -> Self {
262        Self {
263            what: what.into(),
264            location: Location::caller(),
265        }
266    }
267}
268
269impl From<RawException> for Exception {
270    #[track_caller]
271    fn from(e: RawException) -> Self {
272        Self {
273            what: e.what,
274            location: Location::caller(),
275        }
276    }
277}
278
279impl From<&str> for Exception {
280    #[track_caller]
281    fn from(what: &str) -> Self {
282        Self {
283            what: what.to_string(),
284            location: Location::caller(),
285        }
286    }
287}
288
289impl From<Infallible> for Exception {
290    fn from(_: Infallible) -> Self {
291        unreachable!()
292    }
293}
294
295impl From<Exception> for String {
296    fn from(e: Exception) -> Self {
297        e.what
298    }
299}
300
301enum ClosureFailure {
302    Error(Exception),
303    Panic(Box<dyn Any + Send>),
304}
305
306/// Frees an undrained stash when its thread exits; a leaked message would otherwise
307/// outlive every chance to read it.
308struct LastErrorStash(Cell<*const c_char>);
309
310impl LastErrorStash {
311    fn replace(&self, ptr: *const c_char) -> *const c_char {
312        self.0.replace(ptr)
313    }
314}
315
316impl Drop for LastErrorStash {
317    fn drop(&mut self) {
318        let ptr = self.0.replace(std::ptr::null());
319        if !ptr.is_null() {
320            unsafe { libc::free(ptr as *mut libc::c_void) };
321        }
322    }
323}
324
325thread_local! {
326    static CLOSURE_ERROR: Cell<Option<ClosureFailure>> = const { Cell::new(None) };
327    static LAST_MLX_ERROR: LastErrorStash = const { LastErrorStash(Cell::new(std::ptr::null())) };
328}
329
330pub(crate) static INIT_ERR_HANDLER: ErrorHandlerRegistration = ErrorHandlerRegistration::new();
331
332pub(crate) struct ErrorHandlerRegistration(Once);
333
334impl ErrorHandlerRegistration {
335    const fn new() -> Self {
336        Self(Once::new())
337    }
338
339    pub(crate) fn call_once(&self, f: impl FnOnce()) {
340        self.0.call_once(f);
341    }
342
343    pub(crate) fn with<T>(&self, f: impl FnOnce(&Once) -> T) -> T {
344        f(&self.0)
345    }
346}
347
348#[no_mangle]
349extern "C" fn default_mlx_error_handler(msg: *const c_char, _data: *mut std::ffi::c_void) {
350    unsafe {
351        LAST_MLX_ERROR.with(|last_error| {
352            // MLX can report several errors before a caller drains one; keep the
353            // newest message but free the overwritten stash.
354            let previous = last_error.replace(strdup(msg));
355            if !previous.is_null() {
356                libc::free(previous as *mut libc::c_void);
357            }
358        });
359    }
360}
361
362/// Registers one process-global handler; the handler delivers each error through calling-thread
363/// TLS, so one registration serves every thread.
364pub(crate) fn setup_mlx_error_handler() {
365    unsafe {
366        mlx_sys::mlx_set_error_handler(Some(default_mlx_error_handler), std::ptr::null_mut(), None);
367    }
368}
369
370pub(crate) fn set_closure_error(err: Exception) {
371    CLOSURE_ERROR.with(|closure_error| closure_error.set(Some(ClosureFailure::Error(err))));
372}
373
374pub(crate) fn get_and_clear_closure_error() -> Option<Exception> {
375    CLOSURE_ERROR.with(|closure_error| match closure_error.take() {
376        Some(ClosureFailure::Error(err)) => Some(err),
377        Some(ClosureFailure::Panic(payload)) => resume_closure_unwind(payload),
378        None => None,
379    })
380}
381
382pub(crate) fn set_closure_panic(payload: Box<dyn Any + Send>) {
383    CLOSURE_ERROR.with(|closure_error| closure_error.set(Some(ClosureFailure::Panic(payload))));
384}
385
386/// Resumes a panic captured by an MLX closure trampoline after control has returned to Rust.
387///
388/// Closure errors remain available to the calling transform as an [`Exception`]. Panics instead
389/// retain their original payload and resume before the FFI wrapper returns to its caller.
390pub(crate) fn resume_closure_panic() {
391    CLOSURE_ERROR.with(|closure_error| {
392        if let Some(failure) = closure_error.take() {
393            match failure {
394                ClosureFailure::Error(err) => {
395                    closure_error.set(Some(ClosureFailure::Error(err)));
396                }
397                ClosureFailure::Panic(payload) => resume_closure_unwind(payload),
398            }
399        }
400    });
401}
402
403fn resume_closure_unwind(payload: Box<dyn Any + Send>) -> ! {
404    // MLX records its own error when the trampoline reports failure. The panic is the real
405    // cause, so drop that message rather than leaving it to surface against a later operation.
406    let _ = get_and_clear_last_mlx_error();
407    std::panic::resume_unwind(payload)
408}
409
410#[track_caller]
411pub(crate) fn get_and_clear_last_mlx_error() -> Option<RawException> {
412    LAST_MLX_ERROR.with(|last_error| {
413        let last_err_ptr = last_error.replace(std::ptr::null());
414        if last_err_ptr.is_null() {
415            return None;
416        }
417
418        let last_err = unsafe {
419            std::ffi::CStr::from_ptr(last_err_ptr)
420                .to_string_lossy()
421                .into_owned()
422        };
423        unsafe {
424            libc::free(last_err_ptr as *mut libc::c_void);
425        }
426
427        Some(RawException { what: last_err })
428    })
429}
430
431#[track_caller]
432pub(crate) fn exception_from_status(status: i32, operation: &str) -> Exception {
433    get_and_clear_last_mlx_error()
434        .map(Exception::from)
435        .unwrap_or_else(|| Exception::custom(format!("{operation} failed with status {status}")))
436}
437
438/// Error with building a cross-entropy loss function
439#[derive(Debug, Clone, PartialEq, Error)]
440pub enum CrossEntropyBuildError {
441    /// Label smoothing factor must be in the range [0, 1)
442    #[error("Label smoothing factor must be in the range [0, 1)")]
443    InvalidLabelSmoothingFactor,
444}
445
446impl From<CrossEntropyBuildError> for Exception {
447    fn from(value: CrossEntropyBuildError) -> Self {
448        Exception::custom(format!("{value}"))
449    }
450}
451
452/// Error with building a RmsProp optimizer
453#[derive(Debug, Clone, PartialEq, Error)]
454pub enum RmsPropBuildError {
455    /// Alpha must be non-negative
456    #[error("alpha must be non-negative")]
457    NegativeAlpha,
458
459    /// Epsilon must be non-negative
460    #[error("epsilon must be non-negative")]
461    NegativeEpsilon,
462}
463
464/// Error with building an AdaDelta optimizer
465#[derive(Debug, Clone, PartialEq, Error)]
466pub enum AdaDeltaBuildError {
467    /// Rho must be non-negative
468    #[error("rho must be non-negative")]
469    NegativeRho,
470
471    /// Epsilon must be non-negative
472    #[error("epsilon must be non-negative")]
473    NegativeEps,
474}
475
476/// Error with building an Adafactor optimizer.
477#[derive(Debug, Clone, PartialEq, Error)]
478pub enum AdafactorBuildError {
479    /// Either learning rate is provided or relative step is set to true.
480    #[error("Either learning rate is provided or relative step is set to true")]
481    LrIsNoneAndRelativeStepIsFalse,
482}
483
484/// Error with building a dropout layer
485#[derive(Debug, Clone, PartialEq, Error)]
486pub enum DropoutBuildError {
487    /// Dropout probability must be in the range [0, 1)
488    #[error("Dropout probability must be in the range [0, 1)")]
489    InvalidProbability,
490}
491
492/// Error with building a MultiHeadAttention module
493#[derive(Debug, PartialEq, Error)]
494pub enum MultiHeadAttentionBuildError {
495    /// Invalid number of heads
496    #[error("Invalid number of heads: {0}")]
497    InvalidNumHeads(i32),
498
499    /// Exceptions
500    #[error(transparent)]
501    Exception(#[from] Exception),
502}
503
504/// Error with building a transformer
505#[derive(Debug, PartialEq, Error)]
506pub enum TransformerBulidError {
507    /// Dropout probability must be in the range [0, 1)
508    #[error("Dropout probability must be in the range [0, 1)")]
509    InvalidProbability,
510
511    /// Invalid number of heads
512    #[error("Invalid number of heads: {0}")]
513    InvalidNumHeads(i32),
514
515    /// Exceptions
516    #[error(transparent)]
517    Exception(#[from] Exception),
518}
519
520impl From<DropoutBuildError> for TransformerBulidError {
521    fn from(e: DropoutBuildError) -> Self {
522        match e {
523            DropoutBuildError::InvalidProbability => Self::InvalidProbability,
524        }
525    }
526}
527
528impl From<MultiHeadAttentionBuildError> for TransformerBulidError {
529    fn from(e: MultiHeadAttentionBuildError) -> Self {
530        match e {
531            MultiHeadAttentionBuildError::InvalidNumHeads(n) => Self::InvalidNumHeads(n),
532            MultiHeadAttentionBuildError::Exception(e) => Self::Exception(e),
533        }
534    }
535}
536
537/// The dtype is not a float-point type
538#[derive(Debug, Error)]
539#[error("[finfo] dtype {:?} is not inexact", .0)]
540pub struct InexactDtypeError(pub Dtype);
541
542impl From<InexactDtypeError> for Exception {
543    #[track_caller]
544    fn from(value: InexactDtypeError) -> Self {
545        Exception::custom(value.to_string())
546    }
547}
548
549#[cfg(test)]
550mod tests {
551    use crate::array;
552
553    #[test]
554    fn test_exception() {
555        let a = array!([1.0, 2.0, 3.0]);
556        let b = array!([4.0, 5.0]);
557
558        let result = a.add(&b);
559        let error = result.expect_err("Expected error");
560
561        // The full error message would also contain the full path to the original c++ file,
562        // so we just check for a substring
563        assert!(error
564            .what()
565            .contains("Shapes (3) and (2) cannot be broadcast."))
566    }
567}
mlx_rs/error.rs

mlx_rs/
error.rs
