huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-18 19:47:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add some testing for index-add (#237)

Browse Source 
* Add some testing for index-add.

* Fix the cpu implementation for index-add.
This commit is contained in:
Laurent Mazare
2023-07-25 08:38:33 +01:00
committed by GitHub
parent 74a6a769dd
commit 18cc73954a
3 changed files with 66 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
candle-core/src/tensor.rs
View File

@ -1010,18 +1010,26 @@ impl Tensor {
break;
}
}
// The number of element in indexes must match the dimension on which the add is
// performed on the source tensor (and the index values from `indexes` are taken from
// the target tensor self)
mismatch || source_dims[dim] != indexes.dims1()?
mismatch
};
if mismatch {
Err(Error::ShapeMismatchBinaryOp {
op: "index-add",
op: "index-add (self, source)",
lhs: self.shape().clone(),
rhs: source.shape().clone(),
})?
}
// The number of element in indexes must match the dimension on which the add is
// performed on the source tensor (and the index values from `indexes` are taken from
// the target tensor self)
let indexes_len = indexes.dims1()?;
if source_dims[dim] != indexes_len {
Err(Error::ShapeMismatchBinaryOp {
op: "index-add (ids, source))",
lhs: indexes.shape().clone(),
rhs: source.shape().clone(),
})?
}
let storage = self.storage().index_add(
self.layout(),
&indexes.storage(),