huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-17 02:58:50 +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

19
candle-core/src/cpu_backend.rs
View File

@ -815,7 +815,9 @@ impl<'a, I: IntDType> Map2 for IndexAdd<'a, I> {
};
let dim = self.dim;
let max_idx = l1.dims()[dim];
let stride = src_l.stride()[dim];
let pre_dim = src_l.dims()[..dim].iter().product::<usize>();
let src_dim_sz = src_l.dims()[dim];
let post_dim = src_l.dims()[dim + 1..].iter().product::<usize>();
if dim == 0 {
for (src_idx, dst_idx) in self.ids.iter().enumerate() {
let dst_idx = dst_idx.as_usize();
@ -826,17 +828,15 @@ impl<'a, I: IntDType> Map2 for IndexAdd<'a, I> {
size: max_idx,
})?
}
let src_idx = src_idx * stride;
let dst_idx = dst_idx * stride;
let src = &src[src_idx..src_idx + stride];
let dst = &mut dst[dst_idx..dst_idx + stride];
let src_idx = src_idx * post_dim;
let dst_idx = dst_idx * post_dim;
let src = &src[src_idx..src_idx + post_dim];
let dst = &mut dst[dst_idx..dst_idx + post_dim];
for (d, &s) in dst.iter_mut().zip(src.iter()) {
*d += s
}
}
} else {
let pre_dim = src_l.dims()[..dim].iter().product::<usize>();
let post_dim = src_l.dims()[dim + 1..].iter().product::<usize>();
for (src_idx, dst_idx) in self.ids.iter().enumerate() {
let dst_idx = dst_idx.as_usize();
if dst_idx >= max_idx {
@ -847,9 +847,8 @@ impl<'a, I: IntDType> Map2 for IndexAdd<'a, I> {
})?
}
for pre_i in 0..pre_dim {
let pre_i = pre_i * stride;
let pre_src_i = (pre_i + src_idx) * post_dim;
let pre_dst_i = (pre_i + dst_idx) * post_dim;
let pre_src_i = (pre_i * src_dim_sz + src_idx) * post_dim;
let pre_dst_i = (pre_i * max_idx + dst_idx) * post_dim;
let src = &src[pre_src_i..pre_src_i + post_dim];
let dst = &mut dst[pre_dst_i..pre_dst_i + post_dim];
for (d, &s) in dst.iter_mut().zip(src.iter()) {