Skip to main content

hanzo_ml/
error.rs

1//! Hanzo-specific Error and Result
2use std::{convert::Infallible, fmt::Display};
3
4use crate::{DType, DeviceLocation, Layout, MetalError, Shape};
5
6#[derive(Debug, Clone)]
7pub struct MatMulUnexpectedStriding {
8    pub lhs_l: Layout,
9    pub rhs_l: Layout,
10    pub bmnk: (usize, usize, usize, usize),
11    pub msg: &'static str,
12}
13
14impl std::fmt::Debug for Error {
15    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
16        write!(f, "{self}")
17    }
18}
19
20/// Main library error type.
21#[derive(thiserror::Error)]
22pub enum Error {
23    // === DType Errors ===
24    #[error("{msg}, expected: {expected:?}, got: {got:?}")]
25    UnexpectedDType {
26        msg: &'static str,
27        expected: DType,
28        got: DType,
29    },
30
31    #[error("dtype mismatch in {op}, lhs: {lhs:?}, rhs: {rhs:?}")]
32    DTypeMismatchBinaryOp {
33        lhs: DType,
34        rhs: DType,
35        op: &'static str,
36    },
37
38    #[error("unsupported dtype {0:?} for op {1}")]
39    UnsupportedDTypeForOp(DType, &'static str),
40
41    // === Dimension Index Errors ===
42    #[error("{op}: dimension index {dim} out of range for shape {shape:?}")]
43    DimOutOfRange {
44        shape: Shape,
45        dim: i32,
46        op: &'static str,
47    },
48
49    #[error("{op}: duplicate dim index {dims:?} for shape {shape:?}")]
50    DuplicateDimIndex {
51        shape: Shape,
52        dims: Vec<usize>,
53        op: &'static str,
54    },
55
56    // === Shape Errors ===
57    #[error("unexpected rank, expected: {expected}, got: {got} ({shape:?})")]
58    UnexpectedNumberOfDims {
59        expected: usize,
60        got: usize,
61        shape: Shape,
62    },
63
64    #[error("{msg}, expected: {expected:?}, got: {got:?}")]
65    UnexpectedShape {
66        msg: String,
67        expected: Shape,
68        got: Shape,
69    },
70
71    #[error(
72        "Shape mismatch, got buffer of size {buffer_size} which is compatible with shape {shape:?}"
73    )]
74    ShapeMismatch { buffer_size: usize, shape: Shape },
75
76    #[error("shape mismatch in {op}, lhs: {lhs:?}, rhs: {rhs:?}")]
77    ShapeMismatchBinaryOp {
78        lhs: Shape,
79        rhs: Shape,
80        op: &'static str,
81    },
82
83    #[error("shape mismatch in cat for dim {dim}, shape for arg 1: {first_shape:?} shape for arg {n}: {nth_shape:?}")]
84    ShapeMismatchCat {
85        dim: usize,
86        first_shape: Shape,
87        n: usize,
88        nth_shape: Shape,
89    },
90
91    #[error("Cannot divide tensor of shape {shape:?} equally along dim {dim} into {n_parts}")]
92    ShapeMismatchSplit {
93        shape: Shape,
94        dim: usize,
95        n_parts: usize,
96    },
97
98    #[error("{op} can only be performed on a single dimension")]
99    OnlySingleDimension { op: &'static str, dims: Vec<usize> },
100
101    #[error("empty tensor for {op}")]
102    EmptyTensor { op: &'static str },
103
104    // === Device Errors ===
105    #[error("device mismatch in {op}, lhs: {lhs:?}, rhs: {rhs:?}")]
106    DeviceMismatchBinaryOp {
107        lhs: DeviceLocation,
108        rhs: DeviceLocation,
109        op: &'static str,
110    },
111
112    // === Op Specific Errors ===
113    #[error("narrow invalid args {msg}: {shape:?}, dim: {dim}, start: {start}, len:{len}")]
114    NarrowInvalidArgs {
115        shape: Shape,
116        dim: usize,
117        start: usize,
118        len: usize,
119        msg: &'static str,
120    },
121
122    #[error("conv1d invalid args {msg}: inp: {inp_shape:?}, k: {k_shape:?}, pad: {padding}, stride: {stride}")]
123    Conv1dInvalidArgs {
124        inp_shape: Shape,
125        k_shape: Shape,
126        padding: usize,
127        stride: usize,
128        msg: &'static str,
129    },
130
131    #[error("{op} invalid index {index} with dim size {size}")]
132    InvalidIndex {
133        op: &'static str,
134        index: usize,
135        size: usize,
136    },
137
138    #[error("cannot broadcast {src_shape:?} to {dst_shape:?}")]
139    BroadcastIncompatibleShapes { src_shape: Shape, dst_shape: Shape },
140
141    #[error("cannot set variable {msg}")]
142    CannotSetVar { msg: &'static str },
143
144    // Box indirection to avoid large variant.
145    #[error("{0:?}")]
146    MatMulUnexpectedStriding(Box<MatMulUnexpectedStriding>),
147
148    #[error("{op} only supports contiguous tensors")]
149    RequiresContiguous { op: &'static str },
150
151    #[error("{op} expects at least one tensor")]
152    OpRequiresAtLeastOneTensor { op: &'static str },
153
154    #[error("{op} expects at least two tensors")]
155    OpRequiresAtLeastTwoTensors { op: &'static str },
156
157    #[error("backward is not supported for {op}")]
158    BackwardNotSupported { op: &'static str },
159
160    // === Other Errors ===
161    #[error("the hanzo crate has not been built with cuda support")]
162    NotCompiledWithCudaSupport,
163
164    #[error("the hanzo crate has not been built with metal support")]
165    NotCompiledWithMetalSupport,
166
167    #[error("cannot find tensor {path}")]
168    CannotFindTensor { path: String },
169
170    // === Wrapped Errors ===
171    #[error(transparent)]
172    Cuda(Box<dyn std::error::Error + Send + Sync>),
173
174    #[error("Metal error {0}")]
175    Metal(#[from] MetalError),
176
177    #[error(transparent)]
178    TryFromIntError(#[from] core::num::TryFromIntError),
179
180    #[error("npy/npz error {0}")]
181    Npy(String),
182
183    /// Zip file format error.
184    #[error(transparent)]
185    Zip(#[from] zip::result::ZipError),
186
187    /// Integer parse error.
188    #[error(transparent)]
189    ParseInt(#[from] std::num::ParseIntError),
190
191    /// Utf8 parse error.
192    #[error(transparent)]
193    FromUtf8(#[from] std::string::FromUtf8Error),
194
195    /// I/O error.
196    #[error(transparent)]
197    Io(#[from] std::io::Error),
198
199    /// SafeTensor error.
200    #[error(transparent)]
201    SafeTensor(#[from] safetensors::SafeTensorError),
202
203    #[error("unsupported safetensor dtype {0:?}")]
204    UnsupportedSafeTensorDtype(safetensors::Dtype),
205
206    /// Arbitrary errors wrapping.
207    #[error("{0}")]
208    Wrapped(Box<dyn std::fmt::Display + Send + Sync>),
209
210    /// Arbitrary errors wrapping with context.
211    #[error("{wrapped:?}\n{context:?}")]
212    WrappedContext {
213        wrapped: Box<dyn std::error::Error + Send + Sync>,
214        context: String,
215    },
216
217    #[error("{context}\n{inner}")]
218    Context {
219        inner: Box<Self>,
220        context: Box<dyn std::fmt::Display + Send + Sync>,
221    },
222
223    /// Adding path information to an error.
224    #[error("path: {path:?} {inner}")]
225    WithPath {
226        inner: Box<Self>,
227        path: std::path::PathBuf,
228    },
229
230    #[error("{inner}\n{backtrace}")]
231    WithBacktrace {
232        inner: Box<Self>,
233        backtrace: Box<std::backtrace::Backtrace>,
234    },
235
236    /// User generated error message, typically created via `bail!`.
237    #[error("{0}")]
238    Msg(String),
239
240    #[error("unwrap none")]
241    UnwrapNone,
242}
243
244pub type Result<T> = std::result::Result<T, Error>;
245
246impl Error {
247    pub fn wrap(err: impl std::fmt::Display + Send + Sync + 'static) -> Self {
248        Self::Wrapped(Box::new(err)).bt()
249    }
250
251    pub fn msg(err: impl std::fmt::Display) -> Self {
252        Self::Msg(err.to_string()).bt()
253    }
254
255    pub fn debug(err: impl std::fmt::Debug) -> Self {
256        Self::Msg(format!("{err:?}")).bt()
257    }
258
259    pub fn bt(self) -> Self {
260        let backtrace = std::backtrace::Backtrace::capture();
261        match backtrace.status() {
262            std::backtrace::BacktraceStatus::Disabled
263            | std::backtrace::BacktraceStatus::Unsupported => self,
264            _ => Self::WithBacktrace {
265                inner: Box::new(self),
266                backtrace: Box::new(backtrace),
267            },
268        }
269    }
270
271    pub fn with_path<P: AsRef<std::path::Path>>(self, p: P) -> Self {
272        Self::WithPath {
273            inner: Box::new(self),
274            path: p.as_ref().to_path_buf(),
275        }
276    }
277
278    pub fn context(self, c: impl std::fmt::Display + Send + Sync + 'static) -> Self {
279        Self::Context {
280            inner: Box::new(self),
281            context: Box::new(c),
282        }
283    }
284}
285
286#[macro_export]
287macro_rules! bail {
288    ($msg:literal $(,)?) => {
289        return Err($crate::Error::Msg(format!($msg).into()).bt())
290    };
291    ($err:expr $(,)?) => {
292        return Err($crate::Error::Msg(format!($err).into()).bt())
293    };
294    ($fmt:expr, $($arg:tt)*) => {
295        return Err($crate::Error::Msg(format!($fmt, $($arg)*).into()).bt())
296    };
297}
298
299pub fn zip<T, U>(r1: Result<T>, r2: Result<U>) -> Result<(T, U)> {
300    match (r1, r2) {
301        (Ok(r1), Ok(r2)) => Ok((r1, r2)),
302        (Err(e), _) => Err(e),
303        (_, Err(e)) => Err(e),
304    }
305}
306
307pub(crate) mod private {
308    pub trait Sealed {}
309
310    impl<T, E> Sealed for std::result::Result<T, E> where E: std::error::Error {}
311    impl<T> Sealed for Option<T> {}
312}
313
314/// Attach more context to an error.
315///
316/// Inspired by [`anyhow::Context`].
317pub trait Context<T, E>: private::Sealed {
318    /// Wrap the error value with additional context.
319    fn context<C>(self, context: C) -> std::result::Result<T, Error>
320    where
321        C: Display + Send + Sync + 'static;
322
323    /// Wrap the error value with additional context that is evaluated lazily
324    /// only once an error does occur.
325    fn with_context<C, F>(self, f: F) -> std::result::Result<T, Error>
326    where
327        C: Display + Send + Sync + 'static,
328        F: FnOnce() -> C;
329}
330
331impl<T, E> Context<T, E> for std::result::Result<T, E>
332where
333    E: std::error::Error + Send + Sync + 'static,
334{
335    fn context<C>(self, context: C) -> std::result::Result<T, Error>
336    where
337        C: Display + Send + Sync + 'static,
338    {
339        // Not using map_err to save 2 useless frames off the captured backtrace
340        // in ext_context.
341        match self {
342            Ok(ok) => Ok(ok),
343            Err(error) => Err(Error::WrappedContext {
344                wrapped: Box::new(error),
345                context: context.to_string(),
346            }
347            .bt()),
348        }
349    }
350
351    fn with_context<C, F>(self, context: F) -> std::result::Result<T, Error>
352    where
353        C: Display + Send + Sync + 'static,
354        F: FnOnce() -> C,
355    {
356        match self {
357            Ok(ok) => Ok(ok),
358            Err(error) => Err(Error::WrappedContext {
359                wrapped: Box::new(error),
360                context: context().to_string(),
361            }
362            .bt()),
363        }
364    }
365}
366
367impl<T> Context<T, Infallible> for Option<T> {
368    fn context<C>(self, context: C) -> std::result::Result<T, Error>
369    where
370        C: Display + Send + Sync + 'static,
371    {
372        // Not using ok_or_else to save 2 useless frames off the captured
373        // backtrace.
374        match self {
375            Some(ok) => Ok(ok),
376            None => Err(Error::msg(context).bt()),
377        }
378    }
379
380    fn with_context<C, F>(self, context: F) -> std::result::Result<T, Error>
381    where
382        C: Display + Send + Sync + 'static,
383        F: FnOnce() -> C,
384    {
385        match self {
386            Some(v) => Ok(v),
387            None => Err(Error::UnwrapNone.context(context()).bt()),
388        }
389    }
390}