mlx-rs 0.33.0

//! Custom error types and handler for the c ffi

use crate::Dtype;
use libc::strdup;
use std::any::Any;
use std::convert::Infallible;
use std::ffi::NulError;
use std::panic::Location;
use std::sync::Once;
use std::{cell::Cell, ffi::c_char};
use thiserror::Error;

/// Type alias for a `Result` with an `Exception` error type.
pub type Result<T> = std::result::Result<T, Exception>;

/// Error with io operations
#[derive(Error, PartialEq, Debug)]
pub enum IoError {
    /// Path must point to a local file
    #[error("Path must point to a local file")]
    NotFile,

    /// Path contains invalid UTF-8
    #[error("Path contains invalid UTF-8")]
    InvalidUtf8,

    /// Path contains null bytes
    #[error("Path contains null bytes")]
    NullBytes,

    /// No file extension found
    #[error("No file extension found")]
    NoExtension,

    /// Unsupported file format
    #[error("Unsupported file format")]
    UnsupportedFormat,

    /// Unable to open file
    #[error("Unable to open file")]
    UnableToOpenFile,

    /// Unable to allocate memory
    #[error("Unable to allocate memory")]
    AllocationError,

    /// Null error
    #[error(transparent)]
    NulError(#[from] NulError),

    /// Error with unfalttening the loaded optimizer state
    #[error(transparent)]
    Unflatten(#[from] UnflattenError),

    /// Error projecting optimizer state for serialization.
    #[error(transparent)]
    StateProjection(#[from] StateProjectionError),

    /// Exception
    #[error(transparent)]
    Exception(#[from] Exception),
}

impl From<Infallible> for IoError {
    fn from(_: Infallible) -> Self {
        unreachable!()
    }
}

impl From<RawException> for IoError {
    #[track_caller]
    fn from(e: RawException) -> Self {
        let exception = Exception {
            what: e.what,
            location: Location::caller(),
        };
        Self::Exception(exception)
    }
}

/// Error associated with `Array::try_as_slice()`
#[derive(Debug, PartialEq, Error)]
pub enum AsSliceError {
    /// The array data is not contiguous in row-major order.
    #[error(
        "array data is not contiguous row-major; call `Array::contiguous()` before borrowing it as a slice"
    )]
    NotContiguous,

    /// The underlying data pointer is null or the array is empty.
    #[error("The data pointer is null.")]
    Null,

    /// The output dtype does not match the data type of the array.
    #[error("dtype mismatch: expected {expecting:?}, found {found:?}")]
    DtypeMismatch {
        /// The expected data type.
        expecting: Dtype,

        /// The actual data type
        found: Dtype,
    },

    /// Exception
    #[error(transparent)]
    Exception(#[from] Exception),
}

/// Error with unflattening a loaded optimizer state
#[derive(Debug, PartialEq, Error)]
pub enum UnflattenError {
    /// Expecting next (key, value) pair, found none
    #[error("Expecting next (key, value) pair, found none")]
    ExpectingNextPair,

    /// The key is not in a valid format
    #[error("Invalid key")]
    InvalidKey,
}

/// Error with loading an optimizer state
#[derive(Debug, PartialEq, Error)]
pub enum OptimizerStateLoadError {
    /// Error with io operations
    #[error(transparent)]
    Io(#[from] IoError),

    /// Error with unflattening the optimizer state
    #[error(transparent)]
    Unflatten(#[from] UnflattenError),
}

/// Error declaring, restoring, or evaluating a keyed state projection.
#[derive(Debug, PartialEq, Error)]
pub enum StateProjectionError {
    /// A projection declared the same key more than once.
    #[error("duplicate state key {0}")]
    DuplicateKey(String),

    /// A projected key was absent from the supplied snapshot.
    #[error("state snapshot is missing key {0}")]
    MissingKey(String),

    /// The supplied snapshot contains keys absent from the projection.
    #[error("state snapshot contains unknown keys {0:?}")]
    UnknownKeys(Vec<String>),

    /// A required slot was absent in the supplied snapshot.
    #[error("required state key {0} is absent")]
    RequiredSlotAbsent(String),

    /// A compiled state vector has the wrong number of present values.
    #[error("state value count mismatch: expected {expected}, found {actual}")]
    Cardinality {
        /// The number of present values declared by the layout.
        expected: usize,
        /// The supplied value count.
        actual: usize,
    },

    /// The upstream runtime rejected an operation needed to restore state.
    #[error(transparent)]
    Exception(#[from] Exception),
}

impl From<StateProjectionError> for Exception {
    #[track_caller]
    fn from(error: StateProjectionError) -> Self {
        match error {
            StateProjectionError::Exception(error) => error,
            error => Self::custom(error.to_string()),
        }
    }
}

impl From<Infallible> for OptimizerStateLoadError {
    fn from(_: Infallible) -> Self {
        unreachable!()
    }
}

/// Error converting an array or serialized tensor representation.
#[derive(Debug, Error)]
pub enum ConversionError {
    /// The requested element type does not match the array dtype.
    #[error("dtype mismatch: expected {expected:?}, found {actual:?}")]
    DtypeMismatch {
        /// The requested element type.
        expected: Dtype,
        /// The array element type.
        actual: Dtype,
    },

    /// Scalar extraction was requested from a non-scalar array.
    #[error("scalar extraction requires one element, found {actual}")]
    NotScalar {
        /// The array element count.
        actual: usize,
    },

    /// The array cannot be borrowed as a contiguous slice.
    #[error(transparent)]
    ArraySlice(#[from] AsSliceError),

    /// The upstream runtime rejected the conversion.
    #[error(transparent)]
    Exception(#[from] Exception),

    /// The safetensors data type is not supported.
    #[cfg(feature = "safetensors")]
    #[error("The safetensors data type {0:?} is not supported.")]
    SafeTensorDtype(safetensors::tensor::Dtype),

    /// The MLX data type is not supported by safetensors.
    #[cfg(feature = "safetensors")]
    #[error("The mlx data type {0:?} is not supported.")]
    MlxDtype(crate::Dtype),

    /// Error casting the data buffer to `&[u8]`.
    #[cfg(feature = "safetensors")]
    #[error(transparent)]
    PodCastError(#[from] bytemuck::PodCastError),

    /// Error creating a safetensors tensor view.
    #[cfg(feature = "safetensors")]
    #[error(transparent)]
    SafeTensorError(#[from] safetensors::tensor::SafeTensorError),
}

pub(crate) struct RawException {
    pub(crate) what: String,
}

/// Exception. Most will come from the C API.
#[derive(Debug, PartialEq, Error)]
#[error("{what:?} at {location}")]
pub struct Exception {
    pub(crate) what: String,
    pub(crate) location: &'static Location<'static>,
}

impl Exception {
    /// The error message.
    pub fn what(&self) -> &str {
        &self.what
    }

    /// The location of the error.
    ///
    /// The location is obtained from `std::panic::Location::caller()` and points
    /// to the location in the code where the error was created and not where it was
    /// propagated.
    pub fn location(&self) -> &'static Location<'static> {
        self.location
    }

    /// Creates a new exception with the given message.
    #[track_caller]
    pub fn custom(what: impl Into<String>) -> Self {
        Self {
            what: what.into(),
            location: Location::caller(),
        }
    }
}

impl From<RawException> for Exception {
    #[track_caller]
    fn from(e: RawException) -> Self {
        Self {
            what: e.what,
            location: Location::caller(),
        }
    }
}

impl From<&str> for Exception {
    #[track_caller]
    fn from(what: &str) -> Self {
        Self {
            what: what.to_string(),
            location: Location::caller(),
        }
    }
}

impl From<Infallible> for Exception {
    fn from(_: Infallible) -> Self {
        unreachable!()
    }
}

impl From<Exception> for String {
    fn from(e: Exception) -> Self {
        e.what
    }
}

enum ClosureFailure {
    Error(Exception),
    Panic(Box<dyn Any + Send>),
}

/// Frees an undrained stash when its thread exits; a leaked message would otherwise
/// outlive every chance to read it.
struct LastErrorStash(Cell<*const c_char>);

impl LastErrorStash {
    fn replace(&self, ptr: *const c_char) -> *const c_char {
        self.0.replace(ptr)
    }
}

impl Drop for LastErrorStash {
    fn drop(&mut self) {
        let ptr = self.0.replace(std::ptr::null());
        if !ptr.is_null() {
            unsafe { libc::free(ptr as *mut libc::c_void) };
        }
    }
}

thread_local! {
    static CLOSURE_ERROR: Cell<Option<ClosureFailure>> = const { Cell::new(None) };
    static LAST_MLX_ERROR: LastErrorStash = const { LastErrorStash(Cell::new(std::ptr::null())) };
}

pub(crate) static INIT_ERR_HANDLER: ErrorHandlerRegistration = ErrorHandlerRegistration::new();

pub(crate) struct ErrorHandlerRegistration(Once);

impl ErrorHandlerRegistration {
    const fn new() -> Self {
        Self(Once::new())
    }

    pub(crate) fn call_once(&self, f: impl FnOnce()) {
        self.0.call_once(f);
    }

    pub(crate) fn with<T>(&self, f: impl FnOnce(&Once) -> T) -> T {
        f(&self.0)
    }
}

#[no_mangle]
extern "C" fn default_mlx_error_handler(msg: *const c_char, _data: *mut std::ffi::c_void) {
    unsafe {
        LAST_MLX_ERROR.with(|last_error| {
            // MLX can report several errors before a caller drains one; keep the
            // newest message but free the overwritten stash.
            let previous = last_error.replace(strdup(msg));
            if !previous.is_null() {
                libc::free(previous as *mut libc::c_void);
            }
        });
    }
}

/// Registers one process-global handler; the handler delivers each error through calling-thread
/// TLS, so one registration serves every thread.
pub(crate) fn setup_mlx_error_handler() {
    unsafe {
        mlx_sys::mlx_set_error_handler(Some(default_mlx_error_handler), std::ptr::null_mut(), None);
    }
}

pub(crate) fn set_closure_error(err: Exception) {
    CLOSURE_ERROR.with(|closure_error| closure_error.set(Some(ClosureFailure::Error(err))));
}

pub(crate) fn get_and_clear_closure_error() -> Option<Exception> {
    CLOSURE_ERROR.with(|closure_error| match closure_error.take() {
        Some(ClosureFailure::Error(err)) => Some(err),
        Some(ClosureFailure::Panic(payload)) => resume_closure_unwind(payload),
        None => None,
    })
}

pub(crate) fn set_closure_panic(payload: Box<dyn Any + Send>) {
    CLOSURE_ERROR.with(|closure_error| closure_error.set(Some(ClosureFailure::Panic(payload))));
}

/// Resumes a panic captured by an MLX closure trampoline after control has returned to Rust.
///
/// Closure errors remain available to the calling transform as an [`Exception`]. Panics instead
/// retain their original payload and resume before the FFI wrapper returns to its caller.
pub(crate) fn resume_closure_panic() {
    CLOSURE_ERROR.with(|closure_error| {
        if let Some(failure) = closure_error.take() {
            match failure {
                ClosureFailure::Error(err) => {
                    closure_error.set(Some(ClosureFailure::Error(err)));
                }
                ClosureFailure::Panic(payload) => resume_closure_unwind(payload),
            }
        }
    });
}

fn resume_closure_unwind(payload: Box<dyn Any + Send>) -> ! {
    // MLX records its own error when the trampoline reports failure. The panic is the real
    // cause, so drop that message rather than leaving it to surface against a later operation.
    let _ = get_and_clear_last_mlx_error();
    std::panic::resume_unwind(payload)
}

#[track_caller]
pub(crate) fn get_and_clear_last_mlx_error() -> Option<RawException> {
    LAST_MLX_ERROR.with(|last_error| {
        let last_err_ptr = last_error.replace(std::ptr::null());
        if last_err_ptr.is_null() {
            return None;
        }

        let last_err = unsafe {
            std::ffi::CStr::from_ptr(last_err_ptr)
                .to_string_lossy()
                .into_owned()
        };
        unsafe {
            libc::free(last_err_ptr as *mut libc::c_void);
        }

        Some(RawException { what: last_err })
    })
}

#[track_caller]
pub(crate) fn exception_from_status(status: i32, operation: &str) -> Exception {
    get_and_clear_last_mlx_error()
        .map(Exception::from)
        .unwrap_or_else(|| Exception::custom(format!("{operation} failed with status {status}")))
}

/// Error with building a cross-entropy loss function
#[derive(Debug, Clone, PartialEq, Error)]
pub enum CrossEntropyBuildError {
    /// Label smoothing factor must be in the range [0, 1)
    #[error("Label smoothing factor must be in the range [0, 1)")]
    InvalidLabelSmoothingFactor,
}

impl From<CrossEntropyBuildError> for Exception {
    fn from(value: CrossEntropyBuildError) -> Self {
        Exception::custom(format!("{value}"))
    }
}

/// Error with building a RmsProp optimizer
#[derive(Debug, Clone, PartialEq, Error)]
pub enum RmsPropBuildError {
    /// Alpha must be non-negative
    #[error("alpha must be non-negative")]
    NegativeAlpha,

    /// Epsilon must be non-negative
    #[error("epsilon must be non-negative")]
    NegativeEpsilon,
}

/// Error with building an AdaDelta optimizer
#[derive(Debug, Clone, PartialEq, Error)]
pub enum AdaDeltaBuildError {
    /// Rho must be non-negative
    #[error("rho must be non-negative")]
    NegativeRho,

    /// Epsilon must be non-negative
    #[error("epsilon must be non-negative")]
    NegativeEps,
}

/// Error with building an Adafactor optimizer.
#[derive(Debug, Clone, PartialEq, Error)]
pub enum AdafactorBuildError {
    /// Either learning rate is provided or relative step is set to true.
    #[error("Either learning rate is provided or relative step is set to true")]
    LrIsNoneAndRelativeStepIsFalse,
}

/// Error with building a dropout layer
#[derive(Debug, Clone, PartialEq, Error)]
pub enum DropoutBuildError {
    /// Dropout probability must be in the range [0, 1)
    #[error("Dropout probability must be in the range [0, 1)")]
    InvalidProbability,
}

/// Error with building a convolution module
#[derive(Debug, PartialEq, Error)]
pub enum ConvBuildError {
    /// `groups` must be positive and divide both channel counts
    #[error(
        "groups ({groups}) must be positive and divide both input channels ({input_channels}) \
         and output channels ({output_channels})"
    )]
    InvalidGroups {
        /// Requested number of groups
        groups: i32,
        /// Number of input channels
        input_channels: i32,
        /// Number of output channels
        output_channels: i32,
    },

    /// MLX only supports grouped convolution in 1D and 2D
    #[error("grouped 3D convolution is not supported (groups = {0})")]
    GroupedConv3d(i32),

    /// Exceptions
    #[error(transparent)]
    Exception(#[from] Exception),
}

impl From<ConvBuildError> for Exception {
    #[track_caller]
    fn from(value: ConvBuildError) -> Self {
        match value {
            ConvBuildError::Exception(exception) => exception,
            other => Exception::custom(other.to_string()),
        }
    }
}

/// Error with building a MultiHeadAttention module
#[derive(Debug, PartialEq, Error)]
pub enum MultiHeadAttentionBuildError {
    /// Invalid number of heads
    #[error("Invalid number of heads: {0}")]
    InvalidNumHeads(i32),

    /// Exceptions
    #[error(transparent)]
    Exception(#[from] Exception),
}

/// Error with building a transformer
#[derive(Debug, PartialEq, Error)]
pub enum TransformerBulidError {
    /// Dropout probability must be in the range [0, 1)
    #[error("Dropout probability must be in the range [0, 1)")]
    InvalidProbability,

    /// Invalid number of heads
    #[error("Invalid number of heads: {0}")]
    InvalidNumHeads(i32),

    /// Exceptions
    #[error(transparent)]
    Exception(#[from] Exception),
}

impl From<DropoutBuildError> for TransformerBulidError {
    fn from(e: DropoutBuildError) -> Self {
        match e {
            DropoutBuildError::InvalidProbability => Self::InvalidProbability,
        }
    }
}

impl From<MultiHeadAttentionBuildError> for TransformerBulidError {
    fn from(e: MultiHeadAttentionBuildError) -> Self {
        match e {
            MultiHeadAttentionBuildError::InvalidNumHeads(n) => Self::InvalidNumHeads(n),
            MultiHeadAttentionBuildError::Exception(e) => Self::Exception(e),
        }
    }
}

/// The dtype is not a float-point type
#[derive(Debug, Error)]
#[error("[finfo] dtype {:?} is not inexact", .0)]
pub struct InexactDtypeError(pub Dtype);

impl From<InexactDtypeError> for Exception {
    #[track_caller]
    fn from(value: InexactDtypeError) -> Self {
        Exception::custom(value.to_string())
    }
}

#[cfg(test)]
mod tests {
    use crate::array;

    #[test]
    fn test_exception() {
        let a = array!([1.0, 2.0, 3.0]);
        let b = array!([4.0, 5.0]);

        let result = a.add(&b);
        let error = result.expect_err("Expected error");

        // The full error message would also contain the full path to the original c++ file,
        // so we just check for a substring
        assert!(error
            .what()
            .contains("Shapes (3) and (2) cannot be broadcast."))
    }
}