Safetensors changes

dfdx-core/Cargo.toml

@@ -35,7 +35,7 @@ num-traits = { workspace = true }
 safetensors = { workspace = true, optional = true }
 memmap2 = { workspace = true, optional = true }
 half = { version = "2.3.1", optional = true, features = ["num-traits", "rand_distr"] }
-gemm = { version = "0.16.14", default-features = false, optional = true, features = ["rayon"] }
+gemm = { version = "0.17.1", default-features = false, optional = true, features = ["rayon"] }
 rayon = { version = "1.7.0", optional = true }
 libm = { workspace = true }
 wgpu = { version = "0.18.0", features = ["glsl", "spirv"], optional = true }

dfdx-core/src/data/collate.rs

@@ -1,4 +1,4 @@
 use std::{mem::MaybeUninit, vec::Vec};
 
 /// Collates `Self` into some other type.
 /// Generally similar to an unzip method;
@@ -55,6 +55,7 @@
 impl<'a, A, B> Collate for Vec<&'a (A, B)> {
     type Collated = (Vec<&'a A>, Vec<&'a B>);
     fn collated(self) -> Self::Collated {
+        #[allow(clippy::map_identity)]
         self.into_iter().map(|(a, b)| (a, b)).unzip()
     }
 }

dfdx-core/src/lib.rs

@@ -9,7 +9,7 @@
 //! The following sections provide some high level core concepts & exmaples, and
 //! there is more detailed documentation in each of dfdx's submodules.
 //!
 //! See [feature_flags] for details on feature flags.
 //!
 //! # Shapes & Tensors
 //!
@@ -59,7 +59,7 @@
 //! There are two options for this currently, with more planned to be added in the future:
 //!
 //! 1. [tensor::Cpu] - for tensors stored on the heap
 //! 2. [tensor::Cuda] - for tensors stored in GPU memory
 //!
 //! Both devices implement [Default], you can also create them with a certain seed
 //! and ordinal.
@@ -85,8 +85,8 @@
 //! | Unary Operations | `a.sqrt()` | `a.sqrt()` | `a.sqrt()` |
 //! | Binary Operations | `a + b` | `a + b` | `a + b` |
 //! | gemm/gemv | [tensor_ops::matmul] | `a @ b` | `a @ b` |
 //! | 2d Convolution | [tensor_ops::TryConv2D] | - | `torch.conv2d` |
 //! | 2d Transposed Convolution | [tensor_ops::TryConvTrans2D] | - | `torch.conv_transpose2d` |
 //! | Slicing | [tensor_ops::slice] | `a[...]` | `a[...]` |
 //! | Select | [tensor_ops::SelectTo] | `a[...]` | `torch.select` |
 //! | Gather | [tensor_ops::GatherTo] | `np.take` | `torch.gather` |
@@ -128,44 +128,6 @@
     pub use crate::tensor_ops::*;
 }
 
-/// Sets a CPU `sse` flag to flush denormal floating point numbers to zero. The opposite of this is [keep_denormals()].
-///
-/// Some resources:
-/// 1. [Effects of Flush-To-Zero mode](https://developer.arm.com/documentation/dui0473/c/neon-and-vfp-programming/the-effects-of-using-flush-to-zero-mode?lang=en)
-/// 2. [When to use Flush-To-Zero mode](https://developer.arm.com/documentation/dui0473/c/neon-and-vfp-programming/when-to-use-flush-to-zero-mode?lang=en)
-pub fn flush_denormals_to_zero() {
-    #[cfg(all(target_arch = "x86", target_feature = "sse"))]
-    {
-        use std::arch::x86::{_MM_FLUSH_ZERO_ON, _MM_SET_FLUSH_ZERO_MODE};
-        unsafe { _MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_ON) }
-    }
-
-    #[cfg(all(target_arch = "x86_64", target_feature = "sse"))]
-    {
-        use std::arch::x86_64::{_MM_FLUSH_ZERO_ON, _MM_SET_FLUSH_ZERO_MODE};
-        unsafe { _MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_ON) }
-    }
-}
-
-/// Sets a CPU flag to keep denormal floating point numbers. The opposite of this is [flush_denormals_to_zero()].
-///
-/// Some resources:
-/// 1. [Effects of Flush-To-Zero mode](https://developer.arm.com/documentation/dui0473/c/neon-and-vfp-programming/the-effects-of-using-flush-to-zero-mode?lang=en)
-/// 2. [When to use Flush-To-Zero mode](https://developer.arm.com/documentation/dui0473/c/neon-and-vfp-programming/when-to-use-flush-to-zero-mode?lang=en)
-pub fn keep_denormals() {
-    #[cfg(all(target_arch = "x86", target_feature = "sse"))]
-    {
-        use std::arch::x86::{_MM_FLUSH_ZERO_OFF, _MM_SET_FLUSH_ZERO_MODE};
-        unsafe { _MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_OFF) }
-    }
-
-    #[cfg(all(target_arch = "x86_64", target_feature = "sse"))]
-    {
-        use std::arch::x86_64::{_MM_FLUSH_ZERO_OFF, _MM_SET_FLUSH_ZERO_MODE};
-        unsafe { _MM_SET_FLUSH_ZERO_MODE(_MM_FLUSH_ZERO_OFF) }
-    }
-}
-
 #[cfg(test)]
 pub(crate) mod tests {
     pub use num_traits::{Float, NumCast, Zero};

dfdx-core/src/nn_traits/mod.rs

@@ -1,7 +1,7 @@
 mod tuples;
 mod vecs;
 
 use std::vec::Vec;
 
 use crate::prelude::{Device, Dtype, Error, Gradients, Shape, Tensor, UniqueId};
 
@@ -116,12 +116,13 @@
 #[cfg(feature = "safetensors")]
 /// Something that can be saved to a .safetensors file.
 pub trait SaveSafeTensors {
-    fn save_safetensors<P: AsRef<std::path::Path>>(
+    fn save_safetensors_with<P: AsRef<std::path::Path>, F: FnMut(String) -> String>(
         &self,
         path: P,
+        key_map: &mut F,
     ) -> Result<(), safetensors::SafeTensorError> {
         let mut tensors = Vec::new();
-        self.write_safetensors("", &mut tensors);
+        self.write_safetensors_with("", &mut tensors, key_map);
         let data = tensors.iter().map(|(k, dtype, shape, data)| {
             (
                 k.clone(),
@@ -131,53 +132,103 @@
 
         safetensors::serialize_to_file(data, &None, path.as_ref())
     }
-    fn write_safetensors(
+    fn save_safetensors<P: AsRef<std::path::Path>>(
+        &self,
+        path: P,
+    ) -> Result<(), safetensors::SafeTensorError> {
+        self.save_safetensors_with(path, &mut core::convert::identity)
+    }
+    fn write_safetensors_with<F: FnMut(String) -> String>(
         &self,
         location: &str,
         tensors: &mut Vec<(String, safetensors::Dtype, Vec<usize>, Vec<u8>)>,
+        key_map: &mut F,
     );
+    fn write_safetensors(
+        &self,
+        location: &str,
+        tensors: &mut Vec<(String, safetensors::Dtype, Vec<usize>, Vec<u8>)>,
+    ) {
+        self.write_safetensors_with(location, tensors, &mut core::convert::identity)
+    }
 }
 
 #[cfg(feature = "safetensors")]
 /// Something that can be loaded from a .safetensors file.
 pub trait LoadSafeTensors {
-    fn load_safetensors<P: AsRef<std::path::Path>>(
+    fn load_safetensors_with<P: AsRef<std::path::Path>, F: FnMut(String) -> String>(
         &mut self,
         path: P,
+        skip_missing: bool,
+        key_map: &mut F,
     ) -> Result<(), safetensors::SafeTensorError> {
         let f = std::fs::File::open(path)?;
         let buffer = unsafe { memmap2::MmapOptions::new().map(&f)? };
         let tensors = safetensors::SafeTensors::deserialize(&buffer)?;
-        self.read_safetensors("", &tensors)
+        self.read_safetensors_with("", &tensors, skip_missing, key_map)
+    }
+    fn load_safetensors<P: AsRef<std::path::Path>>(
+        &mut self,
+        path: P,
+    ) -> Result<(), safetensors::SafeTensorError> {
+        self.load_safetensors_with(path, false, &mut core::convert::identity)
+    }
+    fn load_safetensors_from_bytes_with<F: FnMut(String) -> String>(
+        &mut self,
+        bytes: &[u8],
+        skip_missing: bool,
+        key_map: &mut F,
+    ) -> Result<(), safetensors::SafeTensorError> {
+        let tensors = safetensors::SafeTensors::deserialize(&bytes)?;
+        self.read_safetensors_with("", &tensors, skip_missing, key_map)
+    }
+    fn load_safetensors_from_bytes(
+        &mut self,
+        bytes: &[u8],
+    ) -> Result<(), safetensors::SafeTensorError> {
+        self.load_safetensors_from_bytes_with(bytes, false, &mut core::convert::identity)
     }
 
-    fn read_safetensors(
+    fn read_safetensors_with<F: FnMut(String) -> String>(
         &mut self,
         location: &str,
         tensors: &safetensors::SafeTensors,
+        skip_missing: bool,
+        key_map: &mut F,
     ) -> Result<(), safetensors::SafeTensorError>;
+    fn read_safetensors(
+        &mut self,
+        location: &str,
+        tensors: &safetensors::SafeTensors,
+    ) -> Result<(), safetensors::SafeTensorError> {
+        self.read_safetensors_with(location, tensors, false, &mut core::convert::identity)
+    }
 }
 
 #[cfg(feature = "safetensors")]
 impl<S: Shape, E: Dtype, D: Device<E>, T> LoadSafeTensors for Tensor<S, E, D, T> {
-    fn read_safetensors(
+    fn read_safetensors_with<F: FnMut(String) -> String>(
         &mut self,
         location: &str,
         tensors: &safetensors::SafeTensors,
+        skip_missing: bool,
+        key_map: &mut F,
     ) -> Result<(), safetensors::SafeTensorError> {
-        self.load_safetensor(tensors, location)
+        self.load_safetensor(tensors, location, skip_missing, key_map)
     }
 }
 
 #[cfg(feature = "safetensors")]
 impl<S: Shape, E: Dtype, D: Device<E>, T> SaveSafeTensors for Tensor<S, E, D, T> {
-    fn write_safetensors(
+    fn write_safetensors_with<F: FnMut(String) -> String>(
         &self,
         location: &str,
         tensors: &mut Vec<(String, safetensors::Dtype, Vec<usize>, Vec<u8>)>,
+        key_map: &mut F,
     ) {
+        let location = key_map(location.to_string());
         tensors.push((
-            location.to_string(),
+            location,
             <E as crate::dtypes::SafeTensorsDtype>::DTYPE,
             self.shape.concrete().into(),
             self.as_vec().iter().flat_map(|e| e.to_le_bytes()).collect(),
@@ -189,15 +240,17 @@
     ($Ty:ty) => {
         #[cfg(feature = "safetensors")]
         impl SaveSafeTensors for $Ty {
-            fn write_safetensors(
+            fn write_safetensors_with<F: FnMut(String) -> String>(
                 &self,
                 location: &str,
                 tensors: &mut Vec<(String, safetensors::Dtype, Vec<usize>, Vec<u8>)>,
+                key_map: &mut F,
             ) {
+                let location = key_map(location.to_string());
                 #[allow(unused_imports)]
                 use crate::dtypes::ToLeBytes;
                 tensors.push((
-                    location.to_string(),
+                    location,
                     <$Ty as crate::dtypes::SafeTensorsDtype>::DTYPE,
                     Vec::new(),
                     self.to_le_bytes().to_vec(),
@@ -207,14 +260,23 @@
 
         #[cfg(feature = "safetensors")]
         impl LoadSafeTensors for $Ty {
-            fn read_safetensors(
+            fn read_safetensors_with<F: FnMut(String) -> String>(
                 &mut self,
                 location: &str,
                 tensors: &safetensors::SafeTensors,
+                skip_missing: bool,
+                key_map: &mut F,
             ) -> Result<(), safetensors::SafeTensorError> {
+                let location = key_map(location.to_string());
                 #[allow(unused_imports)]
                 use crate::dtypes::FromLeBytes;
-                let view = tensors.tensor(location)?;
+                let view = match tensors.tensor(&location) {
+                    Ok(ok) => ok,
+                    Err(safetensors::SafeTensorError::TensorNotFound(_name)) if skip_missing => {
+                        return Ok(());
+                    }
+                    Err(e) => return Err(e),
+                };
                 *self = Self::from_le_bytes(view.data().try_into().unwrap());
                 Ok(())
             }

dfdx-core/src/nn_traits/tuples.rs

@@ -4,7 +4,7 @@
     tensor_ops::Device,
 };
 
 use std::vec::Vec;
 
 macro_rules! tuple_impls {
     ([$($name:ident),+] [$($idx:tt),+], $last:ident, [$($rev_tail:ident),*]) => {
@@ -20,23 +20,32 @@
 
         #[cfg(feature = "safetensors")]
         impl<$($name: crate::nn_traits::SaveSafeTensors, )+> crate::nn_traits::SaveSafeTensors for ($($name,)+) {
-            fn write_safetensors(
+            fn write_safetensors_with<F: FnMut(String) -> String>(
                 &self,
                 location: &str,
                 tensors: &mut Vec<(String, safetensors::Dtype, Vec<usize>, Vec<u8>)>,
+                key_map: &mut F,
             ) {
-                $(self.$idx.write_safetensors(&format!("{location}{}.", $idx), tensors);)+
+                $(
+                    let name = &format!("{location}.{}", $idx);
+                    self.$idx.write_safetensors_with(name, tensors, key_map);
+                )+
             }
         }
 
         #[cfg(feature = "safetensors")]
         impl<$($name: crate::nn_traits::LoadSafeTensors, )+> crate::nn_traits::LoadSafeTensors for ($($name,)+) {
-            fn read_safetensors(
+            fn read_safetensors_with<F: FnMut(String) -> String>(
                 &mut self,
                 location: &str,
                 tensors: &safetensors::SafeTensors,
+                skip_missing: bool,
+                key_map: &mut F,
             ) -> Result<(), safetensors::SafeTensorError> {
-                $(self.$idx.read_safetensors(&format!("{location}{}.", $idx), tensors)?;)+
+                $(
+                    let name = &format!("{location}.{}", $idx);
+                    self.$idx.read_safetensors_with(name, tensors, skip_missing, key_map)?;
+                )+
                 Ok(())
             }
         }

dfdx-core/src/nn_traits/vecs.rs

@@ -4,7 +4,7 @@
     tensor_ops::Device,
 };
 
 use std::vec::Vec;
 
 impl<E: Dtype, D: Device<E>, T: crate::nn_traits::BuildOnDevice<E, D>>
     crate::nn_traits::BuildOnDevice<E, D> for Vec<T>
@@ -60,26 +60,31 @@
 
 #[cfg(feature = "safetensors")]
 impl<T: crate::nn_traits::SaveSafeTensors> crate::nn_traits::SaveSafeTensors for Vec<T> {
-    fn write_safetensors(
+    fn write_safetensors_with<F: FnMut(String) -> String>(
         &self,
         location: &str,
         tensors: &mut Vec<(String, safetensors::Dtype, Vec<usize>, Vec<u8>)>,
+        key_map: &mut F,
     ) {
         for (i, t) in self.iter().enumerate() {
-            t.write_safetensors(&format!("{location}{i}."), tensors);
+            let name = &format!("{location}.{i}");
+            t.write_safetensors_with(name, tensors, key_map);
         }
     }
 }
 
 #[cfg(feature = "safetensors")]
 impl<T: crate::nn_traits::LoadSafeTensors> crate::nn_traits::LoadSafeTensors for Vec<T> {
-    fn read_safetensors(
+    fn read_safetensors_with<F: FnMut(String) -> String>(
         &mut self,
         location: &str,
         tensors: &safetensors::SafeTensors,
+        skip_missing: bool,
+        key_map: &mut F,
     ) -> Result<(), safetensors::SafeTensorError> {
         for (i, t) in self.iter_mut().enumerate() {
-            t.read_safetensors(&format!("{location}{i}."), tensors)?;
+            let name = &format!("{location}.{i}");
+            t.read_safetensors_with(name, tensors, skip_missing, key_map)?;
         }
         Ok(())
     }

dfdx-core/src/tensor/gradients.rs

@@ -153,7 +153,7 @@ impl<E, D: Storage<E>> Gradients<E, D> {
     #[inline]
     pub(crate) fn many_and_ref<L: Shape, R: Shape>(
         &mut self,
-        ls: &Vec<impl Tensorlike<L, E, D>>,
+        ls: &[impl Tensorlike<L, E, D>],
         r: &impl Tensorlike<R, E, D>,
     ) -> (Vec<&mut D::Vec>, &D::Vec) {
         for i in 0..ls.len() {

dfdx-core/src/tensor/mod.rs

@@ -1,10 +1,10 @@
 //! The [Tensor] struct, [Cpu] & [Cuda] devices, and
 //! traits like [ZerosTensor], [OnesTensor], [SampleTensor].
 //!
 //! At a high level a tensor is made up of:
 //! 1. The [crate::shapes::Shape] of the array it stores
 //! 2. The [crate::shapes::Dtype] of the elements of the array
 //! 3. The [Storage] (e.g. [Cpu] or [Cuda]) that it uses to store the nd array
 //! 4. A [Tape], which can either actually be a tape ([OwnedTape])
 //!    or be empty ([NoneTape]).
 //!
@@ -151,7 +151,7 @@
 pub use numpy::NumpyDtype;
 mod error;
 #[cfg(feature = "safetensors")]
-pub mod safetensors;
+mod safetensors;
 mod tensorlike;
 mod unique_id;
 

dfdx-core/src/tensor/safetensors.rs

@@ -5,12 +5,21 @@ use std::vec::Vec;
 
 impl<S: Shape, E: Dtype, D: CopySlice<E>, T> Tensor<S, E, D, T> {
     /// Loads data from the [SafeTensors] `Storage<E>` with the given `key`
-    pub fn load_safetensor(
+    pub fn load_safetensor<F: FnMut(String) -> String>(
         &mut self,
         tensors: &SafeTensors,
         key: &str,
+        skip_missing: bool,
+        key_map: &mut F,
     ) -> Result<(), SafeTensorError> {
-        let tensor_view = tensors.tensor(key)?;
+        let key = key_map(key.to_string());
+        let tensor_view = match tensors.tensor(&key) {
+            Ok(ok) => ok,
+            Err(safetensors::SafeTensorError::TensorNotFound(_name)) if skip_missing => {
+                return Ok(());
+            }
+            Err(e) => return Err(e),
+        };
         let v = tensor_view.data();
         let num_bytes = std::mem::size_of::<E>();
         assert_eq!(