Modularize errors (#4)

Author: walln

Cargo.toml

@@ -5,5 +5,3 @@ members = [
 ]
 
 [workspace.dependencies]
-anyhow = "1.0.75"
-thiserror = "1"

crates/phantom-core/Cargo.toml

@@ -8,5 +8,3 @@ license = "MIT"
 readme = "../../README.md"
 
 [dependencies]
-anyhow = { workspace = true }
-thiserror = { workspace = true }

crates/phantom-core/src/backend/cpu_backend.rs

@@ -1,5 +1,6 @@
 use crate::storage::{BinaryOperation, UnaryOperation};
-use crate::{index::StridedIndex, DType, Error, Result, Shape};
+use crate::{index::StridedIndex, DType, Shape};
+use crate::storage::StorageError;
 
 #[derive(Debug, Clone)]
 pub enum CPUStorage {
@@ -21,7 +22,7 @@ impl CPUStorage {
         stride: &[usize],
         mul: f64,
         add: f64,
-    ) -> Result<Self> {
+    ) -> std::result::Result<Self, StorageError> {
         match self {
             Self::F32(storage) => {
                 let index = StridedIndex::new(shape.dims(), stride);
@@ -42,7 +43,7 @@ impl CPUStorage {
         &self,
         shape: &Shape,
         stride: &[usize],
-    ) -> Result<Self> {
+    ) -> std::result::Result<Self, StorageError> {
         match self {
             Self::F32(storage) => {
                 let index = StridedIndex::new(shape.dims(), stride);
@@ -63,7 +64,7 @@ impl CPUStorage {
         shape: &Shape,
         lhs_stride: &[usize],
         rhs_stride: &[usize],
-    ) -> Result<Self> {
+    ) -> std::result::Result<Self, StorageError> {
         match (self, rhs) {
             (CPUStorage::F32(lhs), CPUStorage::F32(rhs)) => {
                 let lhs_index = StridedIndex::new(shape.dims(), lhs_stride);
@@ -85,16 +86,25 @@ impl CPUStorage {
 
                 Ok(Self::F64(data))
             }
-            _ => Err(Error::BinaryOperationDTypeMismatch {
+            _ => Err(StorageError::BinaryOperationDTypeMismatch {
                 lhs: self.dtype(),
                 rhs: rhs.dtype(),
                 op: T::NAME,
             }),
         }
     }
 
-    pub(crate) fn transpose(&self, shape: &Shape, stride: &[usize]) -> Result<Self> {
-        let (rows, cols) = shape.rank_two()?;
+    pub(crate) fn transpose(&self, shape: &Shape, stride: &[usize]) -> std::result::Result<Self, StorageError> {
+        let (rows, cols) = match shape.rank_two() {
+            Ok(rc) => rc,
+            Err(_) => {
+                return Err(StorageError::BinaryOperationShapeMismatch {
+                    lhs: shape.clone(),
+                    rhs: shape.clone(),
+                    op: "transpose",
+                });
+            }
+        };
         match self {
             CPUStorage::F32(storage) => {
                 let mut out = vec![0f32; rows * cols];
@@ -126,11 +136,11 @@ impl CPUStorage {
         rhs: &Self,
         rhs_shape: (usize, usize),
         rhs_stride: &[usize],
-    ) -> Result<Self> {
+    ) -> std::result::Result<Self, StorageError> {
         let (m, k) = lhs_shape;
         let (k_rhs, n) = rhs_shape;
         if k != k_rhs {
-            return Err(Error::BinaryOperationShapeMismatch {
+            return Err(StorageError::BinaryOperationShapeMismatch {
                 lhs: Shape::from((m, k)),
                 rhs: Shape::from((k_rhs, n)),
                 op: "matmul",
@@ -167,7 +177,7 @@ impl CPUStorage {
                 }
                 Ok(Self::F64(out))
             }
-            _ => Err(Error::BinaryOperationDTypeMismatch {
+            _ => Err(StorageError::BinaryOperationDTypeMismatch {
                 lhs: self.dtype(),
                 rhs: rhs.dtype(),
                 op: "matmul",

crates/phantom-core/src/backprop.rs

@@ -1,6 +1,31 @@
 use crate::tensor::{Tensor, TensorID};
-use crate::Operation;
-use crate::Result;
+use crate::{Operation};
+use crate::tensor::TensorError;
+
+#[derive(Debug)]
+pub enum BackpropError {
+    MissingGradient { tensor: TensorID },
+    Tensor(TensorError),
+}
+
+impl From<TensorError> for BackpropError {
+    fn from(e: TensorError) -> Self {
+        BackpropError::Tensor(e)
+    }
+}
+
+impl std::fmt::Display for BackpropError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            BackpropError::MissingGradient { tensor } => {
+                write!(f, "missing gradient for tensor {:?}", tensor)
+            }
+            BackpropError::Tensor(e) => write!(f, "{e}"),
+        }
+    }
+}
+
+impl std::error::Error for BackpropError {}
 use std::collections::HashMap;
 
 impl Tensor {
@@ -84,9 +109,9 @@ impl Tensor {
     /// let gradients = z.backward()?;
     /// assert_eq!(gradients.len(), 1);
     ///
-    /// # Ok::<(), phantom_core::Error>(())
+    /// # Ok::<(), Box<dyn std::error::Error>>(())
     /// ```
-    pub fn backward(&self) -> Result<HashMap<TensorID, Tensor>> {
+    pub fn backward(&self) -> std::result::Result<HashMap<TensorID, Tensor>, BackpropError> {
         let sorted_nodes = self.sorted_nodes();
         let mut gradients = HashMap::new();
 
@@ -97,7 +122,9 @@ impl Tensor {
                 continue;
             }
 
-            let gradient = gradients.remove(&node.id()).unwrap();
+            let gradient = gradients
+                .remove(&node.id())
+                .ok_or(BackpropError::MissingGradient { tensor: node.id() })?;
 
             if let Some(op) = node.op() {
                 match op {

crates/phantom-core/src/device.rs

@@ -1,5 +1,6 @@
 use crate::backend::cpu_backend::CPUStorage;
-use crate::{storage::Storage, DType, Result, Shape};
+use crate::{storage::Storage, DType, Shape};
+use crate::shape::ShapeError;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum Device {
@@ -41,12 +42,12 @@ impl Device {
 }
 
 pub trait NDArray {
-    fn shape(&self) -> Result<Shape>;
+    fn shape(&self) -> std::result::Result<Shape, ShapeError>;
     fn to_cpu(&self) -> CPUStorage;
 }
 
 impl<S: crate::WithDType> NDArray for S {
-    fn shape(&self) -> Result<Shape> {
+    fn shape(&self) -> std::result::Result<Shape, ShapeError> {
         Ok(Shape::from(()))
     }
 
@@ -56,7 +57,7 @@ impl<S: crate::WithDType> NDArray for S {
 }
 
 impl<S: crate::WithDType> NDArray for &[S] {
-    fn shape(&self) -> Result<Shape> {
+    fn shape(&self) -> std::result::Result<Shape, ShapeError> {
         Ok(Shape::from(self.len()))
     }
 
@@ -66,7 +67,7 @@ impl<S: crate::WithDType> NDArray for &[S] {
 }
 
 impl<S: crate::WithDType, const N: usize> NDArray for &[S; N] {
-    fn shape(&self) -> Result<Shape> {
+    fn shape(&self) -> std::result::Result<Shape, ShapeError> {
         Ok(Shape::from(self.len()))
     }
 
@@ -76,7 +77,7 @@ impl<S: crate::WithDType, const N: usize> NDArray for &[S; N] {
 }
 
 impl<S: crate::WithDType, const N: usize, const M: usize> NDArray for &[[S; N]; M] {
-    fn shape(&self) -> Result<Shape> {
+    fn shape(&self) -> std::result::Result<Shape, ShapeError> {
         Ok(Shape::from((M, N)))
     }
 

crates/phantom-core/src/dtype.rs

@@ -1,4 +1,25 @@
-use crate::{CPUStorage, Error, Result};
+use crate::CPUStorage;
+
+#[derive(Debug)]
+pub enum DTypeError {
+    UnexpectedDType { expected: DType, actual: DType },
+}
+
+impl std::fmt::Display for DTypeError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            DTypeError::UnexpectedDType { expected, actual } => {
+                write!(
+                    f,
+                    "unexpected dtype, expected: {:?}, actual: {:?}",
+                    expected, actual
+                )
+            }
+        }
+    }
+}
+
+impl std::error::Error for DTypeError {}
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
 pub enum DType {
@@ -22,7 +43,7 @@ pub trait WithDType: Sized + Copy {
     fn to_cpu(data: &[Self]) -> CPUStorage {
         Self::to_cpu_owned(data.to_vec())
     }
-    fn storage_slice(storage: &CPUStorage) -> Result<&[Self]>;
+    fn storage_slice(storage: &CPUStorage) -> std::result::Result<&[Self], DTypeError>;
 }
 
 macro_rules! with_dtype {
@@ -34,10 +55,10 @@ macro_rules! with_dtype {
                 CPUStorage::$dtype(data)
             }
 
-            fn storage_slice(storage: &CPUStorage) -> Result<&[Self]> {
+            fn storage_slice(storage: &CPUStorage) -> std::result::Result<&[Self], DTypeError> {
                 match storage {
                     CPUStorage::$dtype(data) => Ok(data),
-                    _ => Err(Error::UnexpectedDType {
+                    _ => Err(DTypeError::UnexpectedDType {
                         expected: DType::$dtype,
                         actual: storage.dtype(),
                     }),

crates/phantom-core/src/error.rs

@@ -2,7 +2,6 @@ mod backend;
 mod backprop;
 mod device;
 mod dtype;
-mod error;
 mod index;
 mod operation;
 mod shape;
@@ -12,9 +11,12 @@ mod tensor;
 pub use backend::cpu_backend::CPUStorage;
 pub use device::Device;
 pub use dtype::{DType, WithDType};
-pub use error::{Error, Result};
+pub use dtype::DTypeError;
 pub use index::StridedIndex;
 pub use operation::Operation;
 pub use shape::Shape;
+pub use shape::ShapeError;
 pub use storage::Storage;
-pub use tensor::{Tensor, TensorID};
+pub use storage::StorageError;
+pub use tensor::{Tensor, TensorError, TensorID};
+pub use backprop::BackpropError;

crates/phantom-core/src/lib.rs

@@ -1,4 +1,20 @@
-use crate::{Error, Result};
+
+#[derive(Debug)]
+pub enum ShapeError {
+    UnexpectedRank { expected: usize, actual: usize, shape: Shape },
+}
+
+impl std::fmt::Display for ShapeError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ShapeError::UnexpectedRank { expected, actual, .. } => {
+                write!(f, "unexpected rank, expected: {}, actual: {}", expected, actual)
+            }
+        }
+    }
+}
+
+impl std::error::Error for ShapeError {}
 
 #[derive(Clone, PartialEq, Eq)]
 pub struct Shape(pub(crate) Vec<usize>);
@@ -47,9 +63,9 @@ impl From<&Shape> for Shape {
 
 macro_rules! get_rank {
     ($fn_name:ident, $cnt:tt, $dims:expr, $out_type:ty) => {
-        pub fn $fn_name(&self) -> Result<$out_type> {
+        pub fn $fn_name(&self) -> std::result::Result<$out_type, ShapeError> {
             if self.0.len() != $cnt {
-                Err(Error::UnexpectedRank {
+                Err(ShapeError::UnexpectedRank {
                     expected: $cnt,
                     actual: self.0.len(),
                     shape: self.clone(),