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

nuvec_core/
error.rs
