Expose candle gather op in pyo3. (#1870)

huggingface · Mar 18, 2024 · 143c481 · 143c481
1 parent f115895
commit 143c481
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/candle-pyo3/py_src/candle/__init__.pyi b/candle-pyo3/py_src/candle/__init__.pyi
@@ -324,6 +324,12 @@ class Tensor:
         """
         pass
 
+    def gather(self, index, dim):
+        """
+        Gathers values along an axis specified by dim.
+        """
+        pass
+
     def get(self, index: int) -> Tensor:
         """
         Gets the value at the specified index.

diff --git a/candle-pyo3/src/lib.rs b/candle-pyo3/src/lib.rs
@@ -448,6 +448,12 @@ impl PyTensor {
         Ok(PyTensor(self.0.index_select(rhs, dim).map_err(wrap_err)?))
     }
 
+    /// Gathers values along an axis specified by dim.
+    fn gather(&self, index: &Self, dim: i64) -> PyResult<Self> {
+        let dim = actual_dim(self, dim).map_err(wrap_err)?;
+        Ok(PyTensor(self.0.gather(index, dim).map_err(wrap_err)?))
+    }
+
     #[pyo3(text_signature = "(self, rhs:Tensor)")]
     /// Performs a matrix multiplication between the two tensors.
     /// &RETURNS&: Tensor