Repository: coreylowman/dfdx (language: rust)

Issue:
Non-tensor attributes of nn objects are not saved by SaveToNpz/LoadFromNpz
This is currently affecting:
- BatchNorm2D::epsilon
- BatchNorm2D::momentum
- Dropout::p
- LayerNorm1D::epsilon


Discussion / hints:
Would it be feasable to store those values as scalar tensor instead and let the `RecursiveWalker` handle the rest upon saving and loading?
Yeah good idea, very straightforward!

Fix patch:
diff --git a/src/nn/batchnorm1d.rs b/src/nn/batchnorm1d.rs
index 848fed421..df0d9fb54 100644
--- a/src/nn/batchnorm1d.rs
+++ b/src/nn/batchnorm1d.rs
@@ -199,14 +199,26 @@ impl<const C: usize, E: Dtype, D: Device<E>> TensorCollection<E, D> for BatchNor
                     |s| &mut s.running_var,
                     TensorOptions::detached(|t| t.try_fill_with_ones()),
                 ),
+                Self::scalar(
+                    "epsilon",
+                    |s| &s.epsilon,
+                    |s| &mut s.epsilon,
+                    ScalarOptions::from_default(1e-5),
+                ),
+                Self::scalar(
+                    "momentum",
+                    |s| &s.momentum,
+                    |s| &mut s.momentum,
+                    ScalarOptions::from_default(0.1),
+                ),
             ),
-            |(scale, bias, running_mean, running_var)| BatchNorm1D {
+            |(scale, bias, running_mean, running_var, epsilon, momentum)| BatchNorm1D {
                 scale,
                 bias,
                 running_mean,
                 running_var,
-                epsilon: 1e-5,
-                momentum: 0.1,
+                epsilon,
+                momentum,
             },
         )
     }
diff --git a/src/nn/batchnorm2d.rs b/src/nn/batchnorm2d.rs
index 4f74b701a..400da5590 100644
--- a/src/nn/batchnorm2d.rs
+++ b/src/nn/batchnorm2d.rs
@@ -272,14 +272,26 @@ impl<const C: usize, E: Dtype, D: Device<E>> TensorCollection<E, D> for BatchNor
                     |s| &mut s.running_var,
                     TensorOptions::detached(|t| t.try_fill_with_ones()),
                 ),
+                Self::scalar(
+                    "epsilon",
+                    |s| &s.epsilon,
+                    |s| &mut s.epsilon,
+                    ScalarOptions::from_default(1e-5),
+                ),
+                Self::scalar(
+                    "momentum",
+                    |s| &s.momentum,
+                    |s| &mut s.momentum,
+                    ScalarOptions::from_default(0.1),
+                ),
             ),
-            |(scale, bias, running_mean, running_var)| BatchNorm2D {
+            |(scale, bias, running_mean, running_var, epsilon, momentum)| BatchNorm2D {
                 scale,
                 bias,
                 running_mean,
                 running_var,
-                epsilon: 1e-5,
-                momentum: 0.1,
+                epsilon,
+                momentum,
             },
         )
     }
diff --git a/src/nn/dropout.rs b/src/nn/dropout.rs
index 018c36d35..7ff12553b 100644
--- a/src/nn/dropout.rs
+++ b/src/nn/dropout.rs
@@ -129,7 +129,27 @@ impl Default for Dropout {
     }
 }
 
-impl ZeroSizedModule for Dropout {}
+impl<D: Device<E>, E: Dtype> BuildOnDevice<D, E> for Dropout {
+    type Built = Dropout;
+}
+
+impl<E: Dtype, D: Device<E>> TensorCollection<E, D> for Dropout {
+    type To<E2: Dtype, D2: Device<E2>> = Dropout;
+
+    fn iter_tensors<V: ModuleVisitor<Self, E, D>>(
+        visitor: &mut V,
+    ) -> Result<Option<Self::To<V::E2, V::D2>>, V::Err> {
+        visitor.visit_fields(
+            <Self as TensorCollection<E, D>>::scalar(
+                "p",
+                |s| &s.p,
+                |s| &mut s.p,
+                ScalarOptions::from_default(0.5),
+            ),
+            |p| Dropout { p },
+        )
+    }
+}
 
 impl<S: Shape, E: Dtype, D: Device<E>> Module<Tensor<S, E, D, NoneTape>> for Dropout {
     type Output = Tensor<S, E, D, NoneTape>;
diff --git a/src/nn/layer_norm.rs b/src/nn/layer_norm.rs
index e867b8704..5d197795e 100644
--- a/src/nn/layer_norm.rs
+++ b/src/nn/layer_norm.rs
@@ -64,11 +64,17 @@ impl<const M: usize, E: Dtype, D: Device<E>> TensorCollection<E, D> for LayerNor
                     |s| &mut s.beta,
                     TensorOptions::reset_to_zeros(),
                 ),
+                Self::scalar(
+                    "epsilon",
+                    |s| &s.epsilon,
+                    |s| &mut s.epsilon,
+                    ScalarOptions::from_default(1e-5),
+                ),
             ),
-            |(gamma, beta)| LayerNorm1D {
+            |(gamma, beta, epsilon)| LayerNorm1D {
                 gamma,
                 beta,
-                epsilon: 1e-5,
+                epsilon,
             },
         )
     }
diff --git a/src/nn/npz.rs b/src/nn/npz.rs
index cc57e63dd..ae0cb574c 100644
--- a/src/nn/npz.rs
+++ b/src/nn/npz.rs
@@ -1,7 +1,7 @@
 use crate::{
     shapes::{Dtype, Shape},
     tensor::{
-        numpy::{NpzError, NumpyDtype},
+        numpy::{read_from_npz, write_to_npz, NpzError, NumpyDtype},
         Tensor,
     },
     tensor_ops::Device,
@@ -129,6 +129,18 @@ impl<W: Write + Seek, E: Dtype + NumpyDtype, D: Device<E>> TensorVisitor<E, D>
         t.write_to_npz(self, full_path)?;
         Ok(None)
     }
+
+    fn visit_scalar<N: num_traits::NumCast>(
+        &mut self,
+        _opts: ScalarOptions<N>,
+        (n, full_path): (&N, String),
+    ) -> Result<Option<N>, Self::Err> {
+        let n = n
+            .to_f64()
+            .unwrap_or_else(|| panic!("Failed to convert scalar value at {full_path} to f64!"));
+        write_to_npz(self, &[], &[n], full_path)?;
+        Ok(None)
+    }
 }
 
 impl<R: Read + Seek, E: Dtype + NumpyDtype, D: Device<E>> TensorVisitor<E, D>
@@ -147,6 +159,22 @@ impl<R: Read + Seek, E: Dtype + NumpyDtype, D: Device<E>> TensorVisitor<E, D>
         t.read_from_npz(self, full_path)?;
         Ok(None)
     }
+
+    fn visit_scalar<N: num_traits::NumCast>(
+        &mut self,
+        _opts: ScalarOptions<N>,
+        (n, full_path): (&mut N, String),
+    ) -> Result<Option<N>, Self::Err> {
+        let buf: Vec<f64> = read_from_npz(self, &[], full_path)?;
+        *n = N::from(buf[0]).unwrap_or_else(|| {
+            panic!(
+                "Failed to convert f64 value {} to {} when reading from npz!",
+                buf[0],
+                std::any::type_name::<N>()
+            )
+        });
+        Ok(None)
+    }
 }
 
 #[cfg(test)]
diff --git a/src/nn/safetensors.rs b/src/nn/safetensors.rs
index f6acf02f7..46f04a80d 100644
--- a/src/nn/safetensors.rs
+++ b/src/nn/safetensors.rs
@@ -12,7 +12,6 @@ use safetensors::{
     tensor::{Dtype as SDtype, SafeTensors, TensorView},
     SafeTensorError,
 };
-use std::collections::BTreeMap;
 
 use super::tensor_collection::*;
 
@@ -25,12 +24,12 @@ struct TensorData {
 }
 
 pub struct Writer {
-    tensors: BTreeMap<String, TensorData>,
+    tensors: Vec<(String, TensorData)>,
 }
 
 impl Writer {
     pub fn new() -> Self {
-        let tensors = BTreeMap::new();
+        let tensors = Vec::new();
         Self { tensors }
     }
 
@@ -44,11 +43,11 @@ impl Writer {
         let data = tensor.as_vec();
         let data: Vec<u8> = data.iter().flat_map(|f| f.to_le_bytes()).collect();
         let tdata = TensorData { dtype, shape, data };
-        self.tensors.insert(key, tdata);
+        self.tensors.push((key, tdata));
     }
 
     pub fn save(&self, path: &Path) -> Result<(), SafeTensorError> {
-        let views: BTreeMap<String, TensorView> = self
+        let (names, views): (Vec<String>, Vec<TensorView>) = self
             .tensors
             .iter()
             .map(|(k, tensor)| {
@@ -57,8 +56,11 @@ impl Writer {
                     TensorView::new(tensor.dtype, tensor.shape.clone(), &tensor.data).unwrap(),
                 )
             })
-            .collect();
-        serialize_to_file(&views, &None, path)
+            .unzip();
+
+        let data = names.into_iter().zip(views.iter());
+
+        serialize_to_file(data, &None, path)
     }
 }
 
@@ -76,6 +78,24 @@ impl<E: Dtype + SafeDtype, D: Device<E>> TensorVisitor<E, D> for Writer {
         self.add(full_path, t);
         Ok(None)
     }
+
+    fn visit_scalar<N: num_traits::NumCast>(
+        &mut self,
+        _: ScalarOptions<N>,
+        (n, full_path): (&N, String),
+    ) -> Result<Option<N>, Self::Err> {
+        let data = TensorData {
+            dtype: safetensors::Dtype::F64,
+            shape: Vec::new(),
+            data: n
+                .to_f64()
+                .unwrap_or_else(|| panic!("Failed to convert scalar value at {full_path} to f64!"))
+                .to_le_bytes()
+                .to_vec(),
+        };
+        self.tensors.push((full_path, data));
+        Ok(None)
+    }
 }
 
 /// Something that can be saved to a `.safetensors`.
@@ -148,6 +168,24 @@ impl<'data, E: Dtype + SafeDtype, D: Device<E>> TensorVisitor<E, D> for SafeTens
         t.load_safetensor(self, &full_path)?;
         Ok(None)
     }
+
+    fn visit_scalar<N: num_traits::NumCast>(
+        &mut self,
+        _: ScalarOptions<N>,
+        (n, full_path): (&mut N, String),
+    ) -> Result<Option<N>, Self::Err> {
+        let data = self.tensor(&full_path)?.data();
+        let mut array = [0; 8];
+        array.copy_from_slice(data);
+        let val = f64::from_le_bytes(array);
+        *n = N::from(val).unwrap_or_else(|| {
+            panic!(
+                "Failed to convert f64 value {val} at {full_path} to {} when reading from safetensors!",
+                std::any::type_name::<N>()
+            )
+        });
+        Ok(None)
+    }
 }
 
 #[cfg(test)]
diff --git a/src/nn/tensor_collection/collection.rs b/src/nn/tensor_collection/collection.rs
index fafe51f8e..8617ffeff 100644
--- a/src/nn/tensor_collection/collection.rs
+++ b/src/nn/tensor_collection/collection.rs
@@ -1,4 +1,5 @@
 #![allow(clippy::type_complexity)]
+use num_traits::NumCast;
 
 use crate::{
     shapes::{ConstShape, Dtype, Shape},
@@ -6,7 +7,7 @@ use crate::{
     tensor_ops::Device,
 };
 
-use super::{ModuleField, ModuleFields, TensorField};
+use super::{ModuleField, ModuleFields, ScalarField, TensorField};
 
 /// A collection of named tensors. Implementing this trait will enable anything
 /// that operates on tensors, including resetting, counting number of params, updating gradients,
@@ -105,6 +106,29 @@ pub trait TensorCollection<E: Dtype, D: Device<E>>: Sized {
             m: Default::default(),
         }
     }
+
+    /// Creates a [ModuleFields] that represents a scalar field.
+    ///
+    /// See also: [TensorField], [TensorCollection], [TensorOptions].
+    fn scalar<F1, F2, N>(
+        name: &str,
+        get_ref: F1,
+        get_mut: F2,
+        options: ScalarOptions<N>,
+    ) -> ScalarField<F1, F2, Self, N>
+    where
+        F1: FnMut(&Self) -> &N,
+        F2: FnMut(&mut Self) -> &mut N,
+        N: NumCast,
+    {
+        ScalarField {
+            name,
+            get_ref,
+            get_mut,
+            options,
+            m: Default::default(),
+        }
+    }
 }
 
 /// An object that can visit [TensorCollection]s and [Tensor]s recursively.
@@ -137,6 +161,18 @@ pub trait ModuleVisitor<T: TensorCollection<E, D>, E: Dtype, D: Device<E>>: Size
         GetRef: FnMut(&T) -> &Tensor<S, E, D>,
         GetMut: FnMut(&mut T) -> &mut Tensor<S, E, D>;
 
+    fn visit_scalar<N, GetRef, GetMut>(
+        &mut self,
+        name: &str,
+        get_refs: GetRef,
+        get_muts: GetMut,
+        opts: ScalarOptions<N>,
+    ) -> Result<Option<N>, Self::Err>
+    where
+        N: NumCast,
+        GetRef: FnMut(&T) -> &N,
+        GetMut: FnMut(&mut T) -> &mut N;
+
     /// Takes something that implements [ModuleFields] and function that takes
     /// [ModuleFields::Output] and returns an instance of T.
     fn visit_fields<M: ModuleFields<T, E, D>>(
@@ -229,3 +265,17 @@ impl<S: Shape, E: Dtype, D: Device<E>> TensorOptions<S, E, D> {
         }
     }
 }
+
+/// Options to change behavior of [ModuleVisitor]
+#[non_exhaustive]
+pub struct ScalarOptions<N: NumCast> {
+    /// The default value for this parameter
+    pub default: N,
+}
+
+impl<N: NumCast> ScalarOptions<N> {
+    // Constructs a ScalarOptions using the parameter's default value
+    pub fn from_default(default: N) -> Self {
+        Self { default }
+    }
+}
diff --git a/src/nn/tensor_collection/mod.rs b/src/nn/tensor_collection/mod.rs
index dbfe10459..b316ea9c4 100644
--- a/src/nn/tensor_collection/mod.rs
+++ b/src/nn/tensor_collection/mod.rs
@@ -6,8 +6,8 @@ mod collection;
 mod visitor;
 mod visitor_impls;
 
-pub use collection::{ModuleVisitor, TensorCollection, TensorOptions};
+pub use collection::{ModuleVisitor, ScalarOptions, TensorCollection, TensorOptions};
 pub use visitor::{
-    ModuleField, ModuleFields, RecursiveWalker, TensorField, TensorViewer, TensorVisitor,
-    ViewTensorMut, ViewTensorName, ViewTensorRef,
+    ModuleField, ModuleFields, RecursiveWalker, ScalarField, TensorField, TensorViewer,
+    TensorVisitor, ViewTensorMut, ViewTensorName, ViewTensorRef,
 };
diff --git a/src/nn/tensor_collection/visitor.rs b/src/nn/tensor_collection/visitor.rs
index 289eff6a7..c8bdbd0e2 100644
--- a/src/nn/tensor_collection/visitor.rs
+++ b/src/nn/tensor_collection/visitor.rs
@@ -1,10 +1,12 @@
+use num_traits::NumCast;
+
 use crate::{
     shapes::{Dtype, Shape},
     tensor::Tensor,
     tensor_ops::Device,
 };
 
-use super::{ModuleVisitor, TensorCollection, TensorOptions};
+use super::{ModuleVisitor, ScalarOptions, TensorCollection, TensorOptions};
 
 /// A standard [ModuleVisitor] that executes `F` on every [Tensor] encountered.
 /// `F` must implement [TensorVisitor]
@@ -77,6 +79,14 @@ pub trait TensorVisitor<E: Dtype, D: Device<E>> {
         opts: TensorOptions<S, E, D>,
         t: <Self::Viewer as TensorViewer>::View<'_, Tensor<S, E, D>>,
     ) -> Result<Option<Tensor<S, Self::E2, Self::D2>>, Self::Err>;
+
+    fn visit_scalar<N: NumCast>(
+        &mut self,
+        opts: ScalarOptions<N>,
+        _h: <Self::Viewer as TensorViewer>::View<'_, N>,
+    ) -> Result<Option<N>, Self::Err> {
+        Ok(Some(opts.default))
+    }
 }
 
 /// Something that can view [Tensor]s in different ways. For example
@@ -110,7 +120,7 @@ pub trait ModuleFields<M: TensorCollection<E, D>, E: Dtype, D: Device<E>> {
     /// and returns optionally constructed fields
     fn visit_fields<V: ModuleVisitor<M, E, D>>(
         self,
-        module: &mut V,
+        visitor: &mut V,
     ) -> Result<Self::Options<V::E2, V::D2>, V::Err>;
 
     /// If any optional fields are None, returns None. Otherwise returns instances of all fields.
@@ -145,6 +155,20 @@ where
     pub(super) m: std::marker::PhantomData<Mod>,
 }
 
+/// A [ModuleFields] that represents a field that contains a scalar value that should be serialized.
+pub struct ScalarField<'a, F1, F2, Mod, N>
+where
+    N: NumCast,
+    F1: FnMut(&Mod) -> &N,
+    F2: FnMut(&mut Mod) -> &mut N,
+{
+    pub(super) name: &'a str,
+    pub(super) get_ref: F1,
+    pub(super) get_mut: F2,
+    pub(super) options: ScalarOptions<N>,
+    pub(super) m: std::marker::PhantomData<Mod>,
+}
+
 /// A [TensorViewer] that represents a `&Tensor`
 #[derive(Debug)]
 pub enum ViewTensorRef {}
diff --git a/src/nn/tensor_collection/visitor_impls.rs b/src/nn/tensor_collection/visitor_impls.rs
index 48612a7f0..4fcb47496 100644
--- a/src/nn/tensor_collection/visitor_impls.rs
+++ b/src/nn/tensor_collection/visitor_impls.rs
@@ -1,3 +1,5 @@
+use num_traits::NumCast;
+
 use std::{
     string::{String, ToString},
     vec::Vec,
@@ -53,6 +55,24 @@ impl<'a, T: TensorCollection<E, D>, E: Dtype, D: Device<E>, F: TensorVisitor<E,
         )
     }
 
+    fn visit_scalar<N, GetRef, GetMut>(
+        &mut self,
+        name: &str,
+        mut get_refs: GetRef,
+        mut get_muts: GetMut,
+        opts: ScalarOptions<N>,
+    ) -> Result<Option<N>, Self::Err>
+    where
+        N: NumCast,
+        GetRef: FnMut(&T) -> &N,
+        GetMut: FnMut(&mut T) -> &mut N,
+    {
+        self.f.visit_scalar(
+            opts,
+            F::Viewer::view_field(&mut self.m, name, &mut get_refs, &mut get_muts),
+        )
+    }
+
     fn visit_fields<M: ModuleFields<T, E, D>>(
         &mut self,
         fields: M,
@@ -219,6 +239,31 @@ where
     }
 }
 
+impl<'a, F1, F2, Mod, N, E: Dtype, D: Device<E>> ModuleFields<Mod, E, D>
+    for ScalarField<'a, F1, F2, Mod, N>
+where
+    N: NumCast,
+    F1: FnMut(&Mod) -> &N,
+    F2: FnMut(&mut Mod) -> &mut N,
+    Mod: TensorCollection<E, D>,
+{
+    type Options<E2: Dtype, D2: Device<E2>> = Option<N>;
+    type Output<E2: Dtype, D2: Device<E2>> = N;
+
+    fn visit_fields<V: ModuleVisitor<Mod, E, D>>(
+        self,
+        visitor: &mut V,
+    ) -> Result<Self::Options<V::E2, V::D2>, V::Err> {
+        visitor.visit_scalar(self.name, self.get_ref, self.get_mut, self.options)
+    }
+
+    fn handle_options<E2: Dtype, D2: Device<E2>>(
+        options: Self::Options<E2, D2>,
+    ) -> Option<Self::Output<E2, D2>> {
+        options
+    }
+}
+
 impl<T: ModuleFields<Mod, E, D>, Mod: TensorCollection<E, D>, E: Dtype, D: Device<E>>
     ModuleFields<Mod, E, D> for Vec<T>
 {
@@ -227,12 +272,12 @@ impl<T: ModuleFields<Mod, E, D>, Mod: TensorCollection<E, D>, E: Dtype, D: Devic
 
     fn visit_fields<V: ModuleVisitor<Mod, E, D>>(
         self,
-        module: &mut V,
+        visitor: &mut V,
     ) -> Result<Self::Options<V::E2, V::D2>, V::Err> {
         let mut out = Vec::with_capacity(self.len());
 
         for x in self {
-            out.push(x.visit_fields(module)?);
+            out.push(x.visit_fields(visitor)?);
         }
 
         Ok(out)
@@ -255,7 +300,7 @@ impl<Mod: TensorCollection<E, D>, E: Dtype, D: Device<E>> ModuleFields<Mod, E, D
     type Options<E2: Dtype, D2: Device<E2>> = ();
     type Output<E2: Dtype, D2: Device<E2>> = ();
 
-    fn visit_fields<V: ModuleVisitor<Mod, E, D>>(self, _module: &mut V) -> Result<(), V::Err> {
+    fn visit_fields<V: ModuleVisitor<Mod, E, D>>(self, _visitor: &mut V) -> Result<(), V::Err> {
         Ok(())
     }
 
@@ -270,7 +315,7 @@ macro_rules! tuple_impls {
             type View<'a, Mod: 'a> = ($($name::View<'a, Mod>,)+);
 
             fn view_field<'a, Mod, Field, GetRef, GetMut>(
-                module: &'a mut Self::View<'_, Mod>,
+                visitor: &'a mut Self::View<'_, Mod>,
                 name: &str,
                 get_ref: &mut GetRef,
                 get_mut: &mut GetMut,
@@ -279,7 +324,7 @@ macro_rules! tuple_impls {
                 GetRef: FnMut(&Mod) -> &Field,
                 GetMut: FnMut(&mut Mod) -> &mut Field,
             {
-                ($($name::view_field(&mut module.$idx, name, get_ref, get_mut),)+)
+                ($($name::view_field(&mut visitor.$idx, name, get_ref, get_mut),)+)
             }
         }
 
diff --git a/src/shapes/shape.rs b/src/shapes/shape.rs
index ec16dbc9e..a0130726b 100644
--- a/src/shapes/shape.rs
+++ b/src/shapes/shape.rs
@@ -70,6 +70,7 @@ pub trait Dtype:
     + std::ops::MulAssign
     + std::ops::DivAssign
     + num_traits::FromPrimitive
+    + num_traits::ToPrimitive
 {
 }
 impl Dtype for f32 {}
diff --git a/src/tensor/numpy.rs b/src/tensor/numpy.rs
index 68f32f0c8..a75404f3d 100644
--- a/src/tensor/numpy.rs
+++ b/src/tensor/numpy.rs
@@ -20,13 +20,10 @@ impl<S: Shape, E: Dtype + NumpyDtype, D: CopySlice<E>, T> Tensor<S, E, D, T> {
     pub fn write_to_npz<W: Write + Seek>(
         &self,
         w: &mut zip::ZipWriter<W>,
-        mut filename: String,
+        filename: String,
     ) -> ZipResult<()> {
-        if !filename.ends_with(".npy") {
-            filename.push_str(".npy");
-        }
-        w.start_file(filename, Default::default())?;
-        self.write_to(w)?;
+        let buf = self.as_vec();
+        write_to_npz(w, self.shape().concrete().as_ref(), &buf, filename)?;
         Ok(())
     }
 
@@ -34,58 +31,92 @@ impl<S: Shape, E: Dtype + NumpyDtype, D: CopySlice<E>, T> Tensor<S, E, D, T> {
     pub fn read_from_npz<R: Read + Seek>(
         &mut self,
         r: &mut zip::ZipArchive<R>,
-        mut filename: String,
+        filename: String,
     ) -> Result<(), NpzError> {
-        if !filename.ends_with(".npy") {
-            filename.push_str(".npy");
-        }
-        let mut f = r
-            .by_name(&filename)
-            .expect(&std::format!("'{}' not found", filename));
-        self.read_from(&mut f)?;
+        let buf = read_from_npz(r, self.shape().concrete().as_ref(), filename)?;
+        self.copy_from(&buf);
         Ok(())
     }
 
     /// Attemps to load the data from a `.npy` file at `path`
     pub fn load_from_npy<P: AsRef<Path>>(&mut self, path: P) -> Result<(), NpyError> {
         let mut f = BufReader::new(File::open(path)?);
-        self.read_from(&mut f)
+        let buf = read_from_npy(&mut f, self.shape().concrete().as_ref())?;
+        self.copy_from(&buf);
+        Ok(())
     }
 
     /// Saves the tensor to a `.npy` file located at `path`
     pub fn save_to_npy<P: AsRef<Path>>(&self, path: P) -> io::Result<()> {
         let mut f = BufWriter::new(File::create(path)?);
-        self.write_to(&mut f)
+        let buf = self.as_vec();
+        write_to_npy(&mut f, self.shape().concrete().as_ref(), &buf)
     }
+}
 
-    pub(crate) fn read_from<R: Read>(&mut self, r: &mut R) -> Result<(), NpyError> {
-        let endian = read_header::<R, E>(r, self.shape().concrete().into_iter().collect())?;
-        let numel = self.shape().num_elements();
-        let mut buf = Vec::with_capacity(numel);
-        for _ in 0..numel {
-            buf.push(E::read_endian(r, endian)?);
-        }
-        D::copy_from(self, &buf);
-        Ok(())
+pub(crate) fn write_to_npz<W: Write + Seek, E: Dtype + NumpyDtype>(
+    w: &mut zip::ZipWriter<W>,
+    shape: &[usize],
+    data: &[E],
+    mut filename: String,
+) -> io::Result<()> {
+    if !filename.ends_with(".npy") {
+        filename.push_str(".npy");
     }
+    w.start_file(filename, Default::default())?;
+    write_to_npy(w, shape, data)
+}
 
-    pub(crate) fn write_to<W: Write>(&self, w: &mut W) -> io::Result<()> {
-        let endian = Endian::Little;
-        write_header::<W, E>(w, endian, self.shape().concrete().into_iter().collect())?;
-        let numel = self.shape().num_elements();
-        let mut buf = std::vec![Default::default(); numel];
-        D::copy_into(self, &mut buf);
-        for v in buf.iter() {
-            v.write_endian(w, endian)?;
-        }
-        Ok(())
+pub(crate) fn read_from_npz<R: Read + Seek, E: Dtype + NumpyDtype>(
+    r: &mut zip::ZipArchive<R>,
+    shape: &[usize],
+    mut filename: String,
+) -> Result<Vec<E>, NpyError> {
+    if !filename.ends_with(".npy") {
+        filename.push_str(".npy");
     }
+
+    let mut f = r
+        .by_name(&filename)
+        .unwrap_or_else(|_| panic!("'{filename}' not found"));
+
+    read_from_npy(&mut f, shape)
+}
+
+fn write_to_npy<W: Write, E: Dtype + NumpyDtype>(
+    w: &mut W,
+    shape: &[usize],
+    data: &[E],
+) -> io::Result<()> {
+    let endian = Endian::Little;
+    write_header::<W, E>(w, endian, shape)?;
+
+    for v in data.iter() {
+        v.write_endian(w, endian)?;
+    }
+
+    Ok(())
+}
+
+fn read_from_npy<R: Read, E: Dtype + NumpyDtype>(
+    r: &mut R,
+    shape: &[usize],
+) -> Result<Vec<E>, NpyError> {
+    let endian = read_header::<R, E>(r, shape)?;
+    let numel = shape.iter().product::<usize>();
+    let mut out = Vec::new();
+
+    for _ in 0..numel {
+        out.push(E::read_endian(r, endian)?);
+    }
+
+    Ok(out)
 }
 
 fn write_header<W: Write, E: NumpyDtype>(
     w: &mut W,
     endian: Endian,
-    shape: Vec<usize>,
+    shape: &[usize],
 ) -> io::Result<()> {
     let shape_str = to_shape_str(shape);
 
@@ -121,7 +152,7 @@ fn write_header<W: Write, E: NumpyDtype>(
     Ok(())
 }
 
-fn read_header<R: Read, E: NumpyDtype>(r: &mut R, shape: Vec<usize>) -> Result<Endian, NpyError> {
+fn read_header<R: Read, E: NumpyDtype>(r: &mut R, shape: &[usize]) -> Result<Endian, NpyError> {
     let mut magic = [0; 6];
     r.read_exact(&mut magic)?;
     if magic != MAGIC_NUMBER {
@@ -311,7 +342,7 @@ impl From<std::string::FromUtf8Error> for NpyError {
     }
 }
 
-fn to_shape_str(shape: Vec<usize>) -> String {
+fn to_shape_str(shape: &[usize]) -> String {
     shape
         .iter()
         .map(|v| v.to_string())
diff --git a/src/tensor/safetensors.rs b/src/tensor/safetensors.rs
index 13a119783..bd717cad9 100644
--- a/src/tensor/safetensors.rs
+++ b/src/tensor/safetensors.rs
@@ -14,12 +14,7 @@ pub trait SafeDtype: Sized {
 impl SafeDtype for f32 {
     type Array = [u8; 4];
     fn from_le_bytes(bytes: &[u8], index: usize) -> Self {
-        Self::from_le_bytes([
-            bytes[index],
-            bytes[index + 1],
-            bytes[index + 2],
-            bytes[index + 3],
-        ])
+        Self::from_le_bytes(bytes[index..index + 4].try_into().unwrap())
     }
 
     fn to_le_bytes(self) -> Self::Array {
@@ -34,16 +29,7 @@ impl SafeDtype for f32 {
 impl SafeDtype for f64 {
     type Array = [u8; 8];
     fn from_le_bytes(bytes: &[u8], index: usize) -> Self {
-        Self::from_le_bytes([
-            bytes[index],
-            bytes[index + 1],
-            bytes[index + 2],
-            bytes[index + 3],
-            bytes[index + 4],
-            bytes[index + 5],
-            bytes[index + 6],
-            bytes[index + 7],
-        ])
+        Self::from_le_bytes(bytes[index..index + 8].try_into().unwrap())
     }
 
     fn safe_dtype() -> SDtype {


Test patch: