nuvec_core/
error.rs

1//! The error type: located op errors (record time) and boundary errors (evaluation, I/O).
2
3use std::fmt;
4use std::panic::Location;
5
6use crate::{DType, Device};
7
8/// Why an operation failed.
9///
10/// Array operations never return this directly: an invalid call yields a *failed array*
11/// that carries an [`Error::Op`] (the op's name, the caller's source location and the
12/// reason), and every array computed from it carries the same error. It surfaces as `Err`
13/// at the boundary: [`crate::Array::check`], [`crate::Array::eval`],
14/// [`crate::Array::item`], [`crate::Array::to_vec`], I/O and transforms. Errors that
15/// happen while evaluating (a GPU failure, memory exhausted) use the other variants.
16#[derive(Debug, Clone, PartialEq, Eq)]
17#[non_exhaustive]
18pub enum Error {
19    /// An array operation rejected its arguments when it was recorded.
20    Op {
21        /// The public operation the caller invoked (MLX's name).
22        op: &'static str,
23        /// Where the caller invoked it.
24        at: &'static Location<'static>,
25        /// Why it failed.
26        kind: OpErrorKind,
27    },
28    /// Memory for this many bytes could not be allocated (or exceeds the memory limit).
29    OutOfMemory(usize),
30    /// The GPU backend failed (for example, a command buffer error).
31    Backend(String),
32    /// No backend is available for the device (the GPU without Metal).
33    UnsupportedDevice(Device),
34    /// Reading or writing a file failed.
35    Io(String),
36    /// An argument of a boundary function (transforms, loaders) is invalid.
37    InvalidArgument(String),
38}
39
40/// The reason an operation rejected its arguments (see [`Error::Op`]).
41#[derive(Debug, Clone, PartialEq, Eq)]
42#[non_exhaustive]
43pub enum OpErrorKind {
44    /// Operand shapes are incompatible.
45    ShapeMismatch {
46        /// Shape of the left operand (or of the array).
47        lhs: Vec<i32>,
48        /// Shape of the right operand (or the requested shape).
49        rhs: Vec<i32>,
50    },
51    /// Operand element types are incompatible.
52    DTypeMismatch {
53        /// Element type of the left operand (or of the array).
54        lhs: DType,
55        /// Element type of the right operand (or the requested type).
56        rhs: DType,
57    },
58    /// The number of host values does not match the shape.
59    SizeMismatch {
60        /// Element count implied by the shape.
61        expected: usize,
62        /// Number of values provided.
63        actual: usize,
64    },
65    /// The operation does not support this element type.
66    UnsupportedDType(DType),
67    /// An axis is out of range, repeated, or more than one where one is expected.
68    InvalidAxis {
69        /// The axis as given.
70        axis: i32,
71        /// Dimensions of the array.
72        ndim: usize,
73    },
74    /// A shape is invalid: a negative or too large dimension, two `-1`, or an element
75    /// count that does not match.
76    InvalidShape(Vec<i64>),
77    /// An index is out of bounds or malformed.
78    InvalidIndex(String),
79    /// Host memory for this many bytes could not be allocated.
80    OutOfMemory(usize),
81    /// An argument is invalid for another reason.
82    InvalidArgument(String),
83}
84
85/// `Result` with [`Error`] as the error type.
86pub type Result<T> = std::result::Result<T, Error>;
87
88/// The location of op errors created inside nuvec, before the public op that serves the
89/// caller attaches the caller's location (see [`Error::locate`]).
90#[inline(never)]
91fn unlocated() -> &'static Location<'static> {
92    Location::caller()
93}
94
95impl Error {
96    /// An op error without a caller location yet: the outermost public op attaches its
97    /// name and the caller's location.
98    pub(crate) fn op(op: &'static str, kind: OpErrorKind) -> Self {
99        Self::Op {
100            op,
101            at: unlocated(),
102            kind,
103        }
104    }
105
106    /// Incompatible operand shapes.
107    pub(crate) fn shape_mismatch(op: &'static str, lhs: &[usize], rhs: &[usize]) -> Self {
108        Self::op(
109            op,
110            OpErrorKind::ShapeMismatch {
111                lhs: dims_i32(lhs),
112                rhs: dims_i32(rhs),
113            },
114        )
115    }
116
117    /// Incompatible operand element types.
118    pub(crate) fn dtype_mismatch(op: &'static str, lhs: DType, rhs: DType) -> Self {
119        Self::op(op, OpErrorKind::DTypeMismatch { lhs, rhs })
120    }
121
122    /// An unsupported element type.
123    pub(crate) fn unsupported_dtype(op: &'static str, dtype: DType) -> Self {
124        Self::op(op, OpErrorKind::UnsupportedDType(dtype))
125    }
126
127    /// Host data that does not match its shape.
128    pub(crate) fn size_mismatch(op: &'static str, expected: usize, actual: usize) -> Self {
129        Self::op(op, OpErrorKind::SizeMismatch { expected, actual })
130    }
131
132    /// An invalid axis.
133    pub(crate) fn invalid_axis(op: &'static str, axis: i64, ndim: usize) -> Self {
134        let axis = i32::try_from(axis).unwrap_or(if axis < 0 { i32::MIN } else { i32::MAX });
135        Self::op(op, OpErrorKind::InvalidAxis { axis, ndim })
136    }
137
138    /// An invalid shape.
139    pub(crate) fn invalid_shape(op: &'static str, shape: impl IntoIterator<Item = i64>) -> Self {
140        Self::op(op, OpErrorKind::InvalidShape(shape.into_iter().collect()))
141    }
142
143    /// An invalid index.
144    pub(crate) fn invalid_index(op: &'static str, msg: impl Into<String>) -> Self {
145        Self::op(op, OpErrorKind::InvalidIndex(msg.into()))
146    }
147
148    /// An op error with the caller's location: `op` (the public operation the caller
149    /// invoked) and `at` replace the internal name of an error that has no location yet;
150    /// invalid arguments and allocation failures become op errors. Errors that already
151    /// carry a location (from a failed input) are returned unchanged: the first error wins.
152    pub(crate) fn locate(self, op: &'static str, at: &'static Location<'static>) -> Self {
153        match self {
154            Self::Op { at: here, kind, .. } if *here == *unlocated() => Self::Op { op, at, kind },
155            Self::InvalidArgument(msg) => Self::Op {
156                op,
157                at,
158                kind: OpErrorKind::InvalidArgument(msg),
159            },
160            Self::OutOfMemory(n) => Self::Op {
161                op,
162                at,
163                kind: OpErrorKind::OutOfMemory(n),
164            },
165            other => other,
166        }
167    }
168
169    /// The reason of an op error.
170    pub fn kind(&self) -> Option<&OpErrorKind> {
171        match self {
172            Self::Op { kind, .. } => Some(kind),
173            _ => None,
174        }
175    }
176
177    /// Where the failing op was called, for op errors raised by a public operation.
178    pub fn location(&self) -> Option<&'static Location<'static>> {
179        match self {
180            Self::Op { at, .. } if **at != *unlocated() => Some(at),
181            _ => None,
182        }
183    }
184
185    /// The name of the failing op, for op errors.
186    pub fn op_name(&self) -> Option<&'static str> {
187        match self {
188            Self::Op { op, .. } => Some(op),
189            _ => None,
190        }
191    }
192}
193
194/// Dimensions as `i32`, saturating (shapes in errors are reported as MLX does).
195pub(crate) fn dims_i32(shape: &[usize]) -> Vec<i32> {
196    shape
197        .iter()
198        .map(|&d| i32::try_from(d).unwrap_or(i32::MAX))
199        .collect()
200}
201
202impl fmt::Display for OpErrorKind {
203    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
204        match self {
205            Self::ShapeMismatch { lhs, rhs } => write!(f, "shape mismatch, {lhs:?} vs {rhs:?}"),
206            Self::DTypeMismatch { lhs, rhs } => write!(f, "dtype mismatch, {lhs} vs {rhs}"),
207            Self::SizeMismatch { expected, actual } => {
208                write!(
209                    f,
210                    "size mismatch, the shape needs {expected} values, got {actual}"
211                )
212            }
213            Self::UnsupportedDType(dtype) => write!(f, "not supported for {dtype}"),
214            Self::InvalidAxis { axis, ndim } => {
215                write!(f, "invalid axis {axis} for an array of {ndim} dimensions")
216            }
217            Self::InvalidShape(shape) => write!(f, "invalid shape {shape:?}"),
218            Self::InvalidIndex(msg) => write!(f, "invalid index: {msg}"),
219            Self::OutOfMemory(n) => write!(f, "out of memory: cannot allocate {n} bytes"),
220            Self::InvalidArgument(msg) => write!(f, "invalid argument: {msg}"),
221        }
222    }
223}
224
225impl fmt::Display for Error {
226    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
227        match self {
228            Self::Op { op, at, kind } if **at == *unlocated() => write!(f, "{op}: {kind}"),
229            Self::Op { op, at, kind } => write!(
230                f,
231                "{op} at {}:{}:{}: {kind}",
232                at.file(),
233                at.line(),
234                at.column()
235            ),
236            Self::OutOfMemory(n) => write!(f, "out of memory: cannot allocate {n} bytes"),
237            Self::Backend(msg) => write!(f, "GPU backend error: {msg}"),
238            Self::UnsupportedDevice(device) => write!(f, "no backend for device {device}"),
239            Self::Io(msg) => write!(f, "I/O error: {msg}"),
240            Self::InvalidArgument(msg) => write!(f, "invalid argument: {msg}"),
241        }
242    }
243}
244
245impl std::error::Error for Error {}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250
251    #[test]
252    fn display_with_location() {
253        let at = Location::caller();
254        let e = Error::shape_mismatch("matmul", &[2, 3], &[4, 5]).locate("matmul", at);
255        assert_eq!(
256            e.to_string(),
257            format!(
258                "matmul at {}:{}:{}: shape mismatch, [2, 3] vs [4, 5]",
259                at.file(),
260                at.line(),
261                at.column()
262            )
263        );
264        assert_eq!(e.location(), Some(at));
265        assert_eq!(e.op_name(), Some("matmul"));
266        assert!(matches!(e.kind(), Some(OpErrorKind::ShapeMismatch { .. })));
267    }
268
269    #[test]
270    fn unlocated_errors_take_the_outer_op() {
271        let inner = Error::unsupported_dtype("exp", DType::Complex64);
272        assert_eq!(inner.location(), None);
273        assert_eq!(inner.to_string(), "exp: not supported for complex64");
274        let at = Location::caller();
275        let outer = inner.locate("softmax", at);
276        assert_eq!(outer.op_name(), Some("softmax"));
277        assert_eq!(outer.location(), Some(at));
278        // Located errors are kept: the first error wins.
279        let again = outer.clone().locate("sum", Location::caller());
280        assert_eq!(again, outer);
281        // Invalid arguments become op errors; boundary errors stay as they are.
282        let e = Error::InvalidArgument("bad".into()).locate("split", at);
283        assert_eq!(e.kind(), Some(&OpErrorKind::InvalidArgument("bad".into())));
284        let io = Error::Io("x".into()).locate("split", at);
285        assert_eq!(io, Error::Io("x".into()));
286    }
287}
nuvec_core/error.rs

nuvec_core/
error.rs
