[Breaking] Removing extra safetensors Error struct

src/nn/safetensors.rs

@@ -1,9 +1,6 @@
 use crate::{
     shapes::{Dtype, HasShape, Shape},
-    tensor::{
-        safetensors::{Error, SafeDtype},
-        CopySlice, Tensor,
-    },
+    tensor::{safetensors::SafeDtype, CopySlice, Tensor},
     tensor_ops::Device,
 };
 use memmap2::MmapOptions;
@@ -136,7 +133,7 @@ pub trait LoadFromSafetensors<E: Dtype + SafeDtype, D: Device<E>>: TensorCollect
     /// let mut model = dev.build_module::<Linear<15, 5>, f32>();
     /// model.load_safetensors("model.safetensors").unwrap();
     /// ```
-    fn load_safetensors<P: AsRef<Path>>(&mut self, path: P) -> Result<(), Error> {
+    fn load_safetensors<P: AsRef<Path>>(&mut self, path: P) -> Result<(), SafeTensorError> {
         let f = std::fs::File::open(path)?;
         let buffer = unsafe { MmapOptions::new().map(&f)? };
         let mut tensors = SafeTensors::deserialize(&buffer)?;
@@ -156,7 +153,7 @@ impl<E: Dtype + SafeDtype, D: Device<E>, T: TensorCollection<E, D>> LoadFromSafe
 
 impl<'data, E: Dtype + SafeDtype, D: Device<E>> TensorVisitor<E, D> for SafeTensors<'data> {
     type Viewer = (ViewTensorMut, ViewTensorName);
-    type Err = Error;
+    type Err = SafeTensorError;
     type E2 = E;
     type D2 = D;
 
@@ -192,7 +189,7 @@ impl<'data, E: Dtype + SafeDtype, D: Device<E>> TensorVisitor<E, D> for SafeTens
                 *n = opts.default;
                 Ok(None)
             }
-            Err(x) => Err(Error::SafeTensorError(x)),
+            Err(x) => Err(x),
         }
     }
 }

src/tensor/safetensors.rs

@@ -11,6 +11,22 @@
     fn safe_dtype() -> SDtype;
 }
 
+#[cfg(feature = "f16")]
+impl SafeDtype for crate::dtypes::f16 {
+    type Array = [u8; 2];
+    fn from_le_bytes(bytes: &[u8], index: usize) -> Self {
+        Self::from_le_bytes(bytes[index..index + 2].try_into().unwrap())
+    }
+
+    fn to_le_bytes(self) -> Self::Array {
+        self.to_le_bytes()
+    }
+
+    fn safe_dtype() -> SDtype {
+        SDtype::F16
+    }
+}
+
 impl SafeDtype for f32 {
     type Array = [u8; 4];
     fn from_le_bytes(bytes: &[u8], index: usize) -> Self {
@@ -41,36 +57,21 @@
     }
 }
 
-#[derive(Debug)]
-pub enum Error {
-    SafeTensorError(SafeTensorError),
-    MismatchedDimension((Vec<usize>, Vec<usize>)),
-    IoError(std::io::Error),
-}
-
-impl From<SafeTensorError> for Error {
-    fn from(safe_error: SafeTensorError) -> Error {
-        Error::SafeTensorError(safe_error)
-    }
-}
-impl From<std::io::Error> for Error {
-    fn from(io_error: std::io::Error) -> Error {
-        Error::IoError(io_error)
-    }
-}
-
 impl<S: Shape, E: Dtype + SafeDtype, D: CopySlice<E>, T> Tensor<S, E, D, T> {
     /// Loads data from the [SafeTensors] Storage<E> with the given `key`
-    pub fn load_safetensor(&mut self, tensors: &SafeTensors, key: &str) -> Result<(), Error> {
-        let tensor = tensors.tensor(key)?;
-        let v = tensor.data();
+    pub fn load_safetensor(
+        &mut self,
+        tensors: &SafeTensors,
+        key: &str,
+    ) -> Result<(), SafeTensorError> {
+        let tensor_view = tensors.tensor(key)?;
+        let v = tensor_view.data();
         let num_bytes = std::mem::size_of::<E>();
-        if tensor.shape() != self.shape.concrete().into() {
-            return Err(Error::MismatchedDimension((
-                tensor.shape().to_vec(),
-                self.shape.concrete().into(),
-            )));
-        }
+        assert_eq!(
+            tensor_view.shape(),
+            self.shape.concrete().into(),
+            "SafeTensors shape did not match tensor shape"
+        );
         if (v.as_ptr() as usize) % num_bytes == 0 {
             // SAFETY This is safe because we just checked that this
             // was correctly aligned.