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

Use the same default as pytorch for sum. (#164)

Browse Source
This commit is contained in:
Laurent Mazare
2023-07-13 21:32:32 +01:00
committed by GitHub
parent 57be3638d8
commit 2bfa791336
13 changed files with 123 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -312,9 +312,11 @@ impl PyTensor {
Ok(PyTensor(self.0.narrow(dim, start, len).map_err(wrap_err)?))
}
fn sum(&self, dims: Vec<usize>) -> PyResult<Self> {
fn sum_keepdim(&self, dims: Vec<usize>) -> PyResult<Self> {
// TODO: Support a single dim as input?
Ok(PyTensor(self.0.sum(dims.as_slice()).map_err(wrap_err)?))
Ok(PyTensor(
self.0.sum_keepdim(dims.as_slice()).map_err(wrap_err)?,
))
}
fn sum_all(&self) -> PyResult<Self> {