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

mlx_rs/
error.rs
