Add varbuilder get_unchecked (huggingface#52)

Author: EricLBuehler

candle-nn/src/var_builder.rs

@@ -1,5 +1,3 @@
-//! A `VarBuilder` for variable retrieval from models
-//!
 //! A `VarBuilder` is used to retrieve variables used by a model. These variables can either come
 //! from a pre-trained checkpoint, e.g. using `VarBuilder::from_mmaped_safetensors`, or initialized
 //! for training, e.g. using `VarBuilder::from_varmap`.
@@ -59,6 +57,9 @@ pub trait Backend: Send + Sync {
         dev: &Device,
     ) -> Result<Tensor>;
 
+    /// Retrieve a tensor based on the name.
+    fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> Result<Tensor>;
+
     fn contains_tensor(&self, name: &str) -> bool;
 }
 
@@ -73,6 +74,9 @@ pub trait SimpleBackend: Send + Sync {
         dev: &Device,
     ) -> Result<Tensor>;
 
+    /// Retrieve a tensor based on the name.
+    fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> Result<Tensor>;
+
     fn contains_tensor(&self, name: &str) -> bool;
 }
 
@@ -89,6 +93,10 @@ impl Backend for Box<dyn SimpleBackend + '_> {
         self.as_ref().get(s, name, h, dtype, dev)
     }
 
+    fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> Result<Tensor> {
+        self.as_ref().get_unchecked(name, dtype, dev)
+    }
+
     fn contains_tensor(&self, name: &str) -> bool {
         self.as_ref().contains_tensor(name)
     }
@@ -194,14 +202,27 @@ impl<B: Backend> VarBuilderArgs<'_, B> {
         name: &str,
         hints: B::Hints,
     ) -> Result<Tensor> {
-        self.get_with_hints_dtype(s, name, hints, self.dtype)
+        self.get_with_hints_dtype(s, name, hints, self.data.dtype)
     }
 
     /// Retrieve the tensor associated with the given name at the current path.
     pub fn get<S: Into<Shape>>(&self, s: S, name: &str) -> Result<Tensor> {
         self.get_with_hints(s, name, Default::default())
     }
 
+    /// Retrieve the tensor associated with the given name at the current path.
+    pub fn get_unchecked(&self, name: &str) -> Result<Tensor> {
+        self.get_unchecked_dtype(name, self.data.dtype)
+    }
+
+    /// Retrieve the tensor associated with the given name & dtype at the current path.
+    pub fn get_unchecked_dtype(&self, name: &str, dtype: DType) -> Result<Tensor> {
+        let name = self.path(name);
+        self.data
+            .backend
+            .get_unchecked(&name, dtype, &self.data.device)
+    }
+
     /// Retrieve the tensor associated with the given name & dtype at the current path.
     pub fn get_with_hints_dtype<S: Into<Shape>>(
         &self,
@@ -224,6 +245,12 @@ impl SimpleBackend for Zeros {
         Tensor::zeros(s, dtype, dev)
     }
 
+    fn get_unchecked(&self, _name: &str, _dtype: DType, _dev: &Device) -> Result<Tensor> {
+        candle::bail!(
+            "`Zeros` requires a shape for tensor retrieval, use `get` instead of `get_unchecked`"
+        )
+    }
+
     fn contains_tensor(&self, _name: &str) -> bool {
         true
     }
@@ -238,6 +265,19 @@ impl SimpleBackend for HashMap<String, Tensor> {
         dtype: DType,
         dev: &Device,
     ) -> Result<Tensor> {
+        let tensor = self.get_unchecked(name, dtype, dev)?;
+        if tensor.shape() != &s {
+            Err(candle::Error::UnexpectedShape {
+                msg: format!("shape mismatch for {name}"),
+                expected: s,
+                got: tensor.shape().clone(),
+            }
+            .bt())?
+        }
+        Ok(tensor)
+    }
+
+    fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> Result<Tensor> {
         let tensor = self
             .get(name)
             .ok_or_else(|| {
@@ -247,14 +287,6 @@ impl SimpleBackend for HashMap<String, Tensor> {
                 .bt()
             })?
             .clone();
-        if tensor.shape() != &s {
-            Err(candle::Error::UnexpectedShape {
-                msg: format!("shape mismatch for {name}"),
-                expected: s,
-                got: tensor.shape().clone(),
-            }
-            .bt())?
-        }
         tensor.to_device(dev)?.to_dtype(dtype)
     }
 
@@ -275,6 +307,10 @@ impl SimpleBackend for VarMap {
         VarMap::get(self, s, name, h, dtype, dev)
     }
 
+    fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> Result<Tensor> {
+        VarMap::get_unchecked(self, name, dtype, dev)
+    }
+
     fn contains_tensor(&self, name: &str) -> bool {
         self.data().lock().unwrap().contains_key(name)
     }
@@ -290,11 +326,24 @@ impl SimpleBackend for SafeTensorWithRouting<'_> {
     fn get(
         &self,
         s: Shape,
-        path: &str,
+        name: &str,
         _: crate::Init,
         dtype: DType,
         dev: &Device,
     ) -> Result<Tensor> {
+        let tensor = self.get_unchecked(name, dtype, dev)?;
+        if tensor.shape() != &s {
+            Err(candle::Error::UnexpectedShape {
+                msg: format!("shape mismatch for {name}"),
+                expected: s,
+                got: tensor.shape().clone(),
+            }
+            .bt())?
+        }
+        Ok(tensor)
+    }
+
+    fn get_unchecked(&self, path: &str, dtype: DType, dev: &Device) -> Result<Tensor> {
         let index = self.routing.get(path).ok_or_else(|| {
             Error::CannotFindTensor {
                 path: path.to_string(),
@@ -305,14 +354,6 @@ impl SimpleBackend for SafeTensorWithRouting<'_> {
             .tensor(path)?
             .load(dev)?
             .to_dtype(dtype)?;
-        if tensor.shape() != &s {
-            Err(candle::Error::UnexpectedShape {
-                msg: format!("shape mismatch for {path}"),
-                expected: s,
-                got: tensor.shape().clone(),
-            }
-            .bt())?
-        }
         Ok(tensor)
     }
 
@@ -325,22 +366,15 @@ impl SimpleBackend for candle::npy::NpzTensors {
     fn get(
         &self,
         s: Shape,
-        path: &str,
+        name: &str,
         _: crate::Init,
         dtype: DType,
         dev: &Device,
     ) -> Result<Tensor> {
-        let tensor = match self.get(path)? {
-            None => Err(Error::CannotFindTensor {
-                path: path.to_string(),
-            }
-            .bt())?,
-            Some(tensor) => tensor,
-        };
-        let tensor = tensor.to_device(dev)?.to_dtype(dtype)?;
+        let tensor = self.get_unchecked(name, dtype, dev)?;
         if tensor.shape() != &s {
             Err(candle::Error::UnexpectedShape {
-                msg: format!("shape mismatch for {path}"),
+                msg: format!("shape mismatch for {name}"),
                 expected: s,
                 got: tensor.shape().clone(),
             }
@@ -349,6 +383,18 @@ impl SimpleBackend for candle::npy::NpzTensors {
         Ok(tensor)
     }
 
+    fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> Result<Tensor> {
+        let tensor = match self.get(name)? {
+            None => Err(Error::CannotFindTensor {
+                path: name.to_string(),
+            }
+            .bt())?,
+            Some(tensor) => tensor,
+        };
+        let tensor = tensor.to_device(dev)?.to_dtype(dtype)?;
+        Ok(tensor)
+    }
+
     fn contains_tensor(&self, name: &str) -> bool {
         self.get(name).is_ok_and(|v| v.is_some())
     }
@@ -358,22 +404,15 @@ impl SimpleBackend for candle::pickle::PthTensors {
     fn get(
         &self,
         s: Shape,
-        path: &str,
+        name: &str,
         _: crate::Init,
         dtype: DType,
         dev: &Device,
     ) -> Result<Tensor> {
-        let tensor = match self.get(path)? {
-            None => Err(Error::CannotFindTensor {
-                path: path.to_string(),
-            }
-            .bt())?,
-            Some(tensor) => tensor,
-        };
-        let tensor = tensor.to_device(dev)?.to_dtype(dtype)?;
+        let tensor = self.get_unchecked(name, dtype, dev)?;
         if tensor.shape() != &s {
             Err(candle::Error::UnexpectedShape {
-                msg: format!("shape mismatch for {path}"),
+                msg: format!("shape mismatch for {name}"),
                 expected: s,
                 got: tensor.shape().clone(),
             }
@@ -382,6 +421,18 @@ impl SimpleBackend for candle::pickle::PthTensors {
         Ok(tensor)
     }
 
+    fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> Result<Tensor> {
+        let tensor = match self.get(name)? {
+            None => Err(Error::CannotFindTensor {
+                path: name.to_string(),
+            }
+            .bt())?,
+            Some(tensor) => tensor,
+        };
+        let tensor = tensor.to_device(dev)?.to_dtype(dtype)?;
+        Ok(tensor)
+    }
+
     fn contains_tensor(&self, name: &str) -> bool {
         self.get(name).is_ok_and(|v| v.is_some())
     }
@@ -396,7 +447,7 @@ impl SimpleBackend for candle::safetensors::MmapedSafetensors {
         dtype: DType,
         dev: &Device,
     ) -> Result<Tensor> {
-        let tensor = self.load(name, dev)?.to_dtype(dtype)?;
+        let tensor = self.get_unchecked(name, dtype, dev)?;
         if tensor.shape() != &s {
             Err(candle::Error::UnexpectedShape {
                 msg: format!("shape mismatch for {name}"),
@@ -408,6 +459,10 @@ impl SimpleBackend for candle::safetensors::MmapedSafetensors {
         Ok(tensor)
     }
 
+    fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> Result<Tensor> {
+        self.load(name, dev)?.to_dtype(dtype)
+    }
+
     fn contains_tensor(&self, name: &str) -> bool {
         self.get(name).is_ok()
     }
@@ -422,7 +477,7 @@ impl SimpleBackend for candle::safetensors::BufferedSafetensors {
         dtype: DType,
         dev: &Device,
     ) -> Result<Tensor> {
-        let tensor = self.load(name, dev)?.to_dtype(dtype)?;
+        let tensor = self.get_unchecked(name, dtype, dev)?;
         if tensor.shape() != &s {
             Err(candle::Error::UnexpectedShape {
                 msg: format!("shape mismatch for {name}"),
@@ -434,6 +489,10 @@ impl SimpleBackend for candle::safetensors::BufferedSafetensors {
         Ok(tensor)
     }
 
+    fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> Result<Tensor> {
+        self.load(name, dev)?.to_dtype(dtype)
+    }
+
     fn contains_tensor(&self, name: &str) -> bool {
         self.get(name).is_ok()
     }
@@ -448,7 +507,7 @@ impl SimpleBackend for candle::safetensors::SliceSafetensors<'_> {
         dtype: DType,
         dev: &Device,
     ) -> Result<Tensor> {
-        let tensor = self.load(name, dev)?.to_dtype(dtype)?;
+        let tensor = self.get_unchecked(name, dtype, dev)?;
         if tensor.shape() != &s {
             Err(candle::Error::UnexpectedShape {
                 msg: format!("shape mismatch for {name}"),
@@ -460,6 +519,10 @@ impl SimpleBackend for candle::safetensors::SliceSafetensors<'_> {
         Ok(tensor)
     }
 
+    fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> Result<Tensor> {
+        self.load(name, dev)?.to_dtype(dtype)
+    }
+
     fn contains_tensor(&self, name: &str) -> bool {
         self.get(name).is_ok()
     }
@@ -714,6 +777,10 @@ impl Backend for ShardedSafeTensors {
         Tensor::from_raw_buffer(&raw, view_dtype, &shape, dev)?.to_dtype(dtype)
     }
 
+    fn get_unchecked(&self, _name: &str, _dtype: DType, _dev: &Device) -> Result<Tensor> {
+        candle::bail!("`get_unchecked` does not make sense for `ShardedSafeTensors`, use `get`.");
+    }
+
     fn contains_tensor(&self, name: &str) -> bool {
         self.0.get(name).is_ok()
     }
@@ -747,6 +814,11 @@ impl<R: Renamer + Sync + Send> SimpleBackend for Rename<'_, R> {
             .to_device(dev)
     }
 
+    fn get_unchecked(&self, name: &str, dtype: DType, dev: &Device) -> Result<Tensor> {
+        let name = self.renamer.rename(name);
+        self.inner.get_unchecked_dtype(&name, dtype)?.to_device(dev)
+    }
+
     fn contains_tensor(&self, name: &str) -> bool {
         let name = self.renamer.rename(name);
         self.inner.contains_tensor(&name)

candle-nn/src/var_map.rs

@@ -1,5 +1,3 @@
-//! A `VarMap` is a store that holds named variables.
-//!
 use candle::{DType, Device, Result, Shape, Tensor, Var};
 use std::collections::HashMap;
 use std::sync::{Arc, Mutex};
@@ -115,6 +113,11 @@ impl VarMap {
         Ok(tensor)
     }
 
+    /// Retrieve or add a new variable.
+    pub fn get_unchecked(&self, _path: &str, _dtype: DType, _device: &Device) -> Result<Tensor> {
+        candle::bail!("`get_unchecked` does not make sense for `VarMap`, use `get`.");
+    }
+
     pub fn data(&self) -> &Mutex<HashMap<String, Var>> {
         &self.data
     }