huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-15 02:16:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Expose candle gather op in pyo3. (#1870)

Browse Source
This commit is contained in:
Laurent Mazare
2024-03-18 21:54:15 +01:00
committed by GitHub
parent f115895b9e
commit 143c481c20
2 changed files with 12 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
candle-pyo3/py_src/candle/__init__.pyi
View File

@ -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.

6
candle-pyo3/src/lib.rs
View File

@ -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