Update safetensors module and naming

- Makes the safetensors module private. - Doesn't get exported on the preamble, avoiding a naming clash with the safetensors external crate. - Change how and when the period is inserted. - This should make it closer to how the fields are accessed in the code.
coreylowman · Feb 1, 2024 · a832f51 · a832f51
1 parent 4722a99
commit a832f51
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 9 deletions.
diff --git a/dfdx-core/src/nn_traits/tuples.rs b/dfdx-core/src/nn_traits/tuples.rs
@@ -25,7 +25,7 @@ macro_rules! tuple_impls {
                 location: &str,
                 tensors: &mut Vec<(String, safetensors::Dtype, Vec<usize>, Vec<u8>)>,
             ) {
-                $(self.$idx.write_safetensors(&format!("{location}{}.", $idx), tensors);)+
+                $(self.$idx.write_safetensors(&format!("{location}.{}", $idx), tensors);)+
             }
         }
 
@@ -36,7 +36,7 @@ macro_rules! tuple_impls {
                 location: &str,
                 tensors: &safetensors::SafeTensors,
             ) -> Result<(), safetensors::SafeTensorError> {
-                $(self.$idx.read_safetensors(&format!("{location}{}.", $idx), tensors)?;)+
+                $(self.$idx.read_safetensors(&format!("{location}.{}", $idx), tensors)?;)+
                 Ok(())
             }
         }

diff --git a/dfdx-core/src/nn_traits/vecs.rs b/dfdx-core/src/nn_traits/vecs.rs
@@ -66,7 +66,7 @@ impl<T: crate::nn_traits::SaveSafeTensors> crate::nn_traits::SaveSafeTensors for
         tensors: &mut Vec<(String, safetensors::Dtype, Vec<usize>, Vec<u8>)>,
     ) {
         for (i, t) in self.iter().enumerate() {
-            t.write_safetensors(&format!("{location}{i}."), tensors);
+            t.write_safetensors(&format!("{location}.{i}"), tensors);
         }
     }
 }
@@ -79,7 +79,7 @@ impl<T: crate::nn_traits::LoadSafeTensors> crate::nn_traits::LoadSafeTensors for
         tensors: &safetensors::SafeTensors,
     ) -> Result<(), safetensors::SafeTensorError> {
         for (i, t) in self.iter_mut().enumerate() {
-            t.read_safetensors(&format!("{location}{i}."), tensors)?;
+            t.read_safetensors(&format!("{location}.{i}"), tensors)?;
         }
         Ok(())
     }

diff --git a/dfdx-core/src/tensor/mod.rs b/dfdx-core/src/tensor/mod.rs
@@ -151,7 +151,7 @@ pub(crate) mod webgpu;
 pub use numpy::NumpyDtype;
 mod error;
 #[cfg(feature = "safetensors")]
-pub mod safetensors;
+mod safetensors;
 mod tensorlike;
 mod unique_id;
 

diff --git a/dfdx-derives/src/lib.rs b/dfdx-derives/src/lib.rs
@@ -850,7 +850,10 @@ pub fn save_safetensors(input: proc_macro::TokenStream) -> proc_macro::TokenStre
                         where_clause
                             .predicates
                             .push(parse_quote!(#ty: ::dfdx::nn_traits::SaveSafeTensors));
-                        quote_spanned!(f.span()=>self.#name.write_safetensors(&format!("{location}{}", #name_str), tensors);)
+                        quote_spanned!(f.span()=>self.#name.write_safetensors(
+                            &format!("{location}{}{}", if location.is_empty() { "" } else { "." },  #name_str), 
+                            tensors
+                        );)
                     } else {
                         Default::default()
                     }
@@ -866,7 +869,10 @@ pub fn save_safetensors(input: proc_macro::TokenStream) -> proc_macro::TokenStre
                         where_clause
                             .predicates
                             .push(parse_quote!(#ty: ::dfdx::nn_traits::SaveSafeTensors));
-                        quote_spanned!(f.span()=>self.#index.write_safetensors(&format!("{location}{}", #index), tensors);)
+                        quote_spanned!(f.span()=>self.#index.write_safetensors(
+                            &format!("{location}{}{}", if location.is_empty() { "" } else { "." }, #index), 
+                            tensors
+                        );)
                     } else {
                         Default::default()
                     }
@@ -913,7 +919,10 @@ pub fn load_safetensors(input: proc_macro::TokenStream) -> proc_macro::TokenStre
                         where_clause
                             .predicates
                             .push(parse_quote!(#ty: ::dfdx::nn_traits::LoadSafeTensors));
-                        quote_spanned!(f.span()=>self.#name.read_safetensors(&format!("{location}{}", #name_str), tensors)?;)
+                        quote_spanned!(f.span()=>self.#name.read_safetensors(
+                            &format!("{location}{}{}", if location.is_empty() { "" } else { "." }, #name_str),
+                            tensors
+                        )?;)
                     } else {
                         Default::default()
                     }
@@ -928,7 +937,10 @@ pub fn load_safetensors(input: proc_macro::TokenStream) -> proc_macro::TokenStre
                         where_clause
                             .predicates
                             .push(parse_quote!(#ty: ::dfdx::nn_traits::LoadSafeTensors));
-                        quote_spanned!(f.span()=>self.#index.read_safetensors(&format!("{location}{}", #index), tensors)?;)
+                        quote_spanned!(f.span()=>self.#index.read_safetensors(
+                            &format!("{location}{}{}", if location.is_empty() { "" } else { "." }, #index),
+                            tensors
+                        )?;)
                     } else {
                         Default::default()
                     }