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

Fix the reduce_sum implementation and add some tests.

Browse Source
This commit is contained in:
laurent
2023-06-25 10:55:04 +01:00
parent 0f369dd870
commit ba0693a908
2 changed files with 29 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/cpu_backend.rs
View File

@ -178,18 +178,22 @@ impl CpuStorage {
dst_dims[sum_dim] = 1;
}
let dst_shape = Shape::from(dst_dims);
let sum_dims_and_stride: Vec<_> = src_dims
let mut sum_dims = sum_dims.to_vec();
// Sort the sum_dims as they have to be processed from left to right when converting the
// indexes.
sum_dims.sort();
let sum_dims_and_stride: Vec<_> = sum_dims
.iter()
.enumerate()
.map(|(i, d)| (d, src_dims[i + 1..].iter().product::<usize>()))
.map(|&d| (src_dims[d], src_dims[d + 1..].iter().product::<usize>()))
.collect();
let to_dst_index = |unstr_index: usize| {
// TODO: Optimize, the following does lots of slow division and modulos.
// TODO: Optimize, the following does lots of slow division.
let mut dst_index = unstr_index;
// Set the sum_dims indexes to 0.
for &(dim, stride) in sum_dims_and_stride.iter() {
let index = dst_index / stride % dim;
dst_index -= index * stride;
// The compiler is able to optimize the following in a single divmod op.
let (pre, post) = (dst_index / stride, dst_index % stride);
dst_index = (pre / dim) * stride + post;
}
dst_index
};

19
tests/tensor_tests.rs
View File

@ -108,6 +108,25 @@ fn softmax() -> Result<()> {
Ok(())
}
#[test]
fn sum() -> Result<()> {
let data = &[[[3u32, 1, 4], [1, 5, 9]], [[2, 1, 7], [8, 2, 8]]];
let tensor = Tensor::new(data, &Device::Cpu)?;
assert_eq!(
tensor.sum(&[2])?.to_vec3::<u32>()?,
&[[[8], [15]], [[10], [18]]]
);
assert_eq!(
tensor.t()?.sum(&[1])?.t()?.to_vec3::<u32>()?,
&[[[8], [15]], [[10], [18]]]
);
assert_eq!(
tensor.sum(&[2, 1])?.to_vec3::<u32>()?,
&[[[8 + 15]], [[10 + 18]]]
);
Ok(())
}
#[test]
fn narrow() -> Result<()> {
let data = &[[[3f32, 1., 4.], [1., 5., 9.]], [[2., 1., 7.], [8., 2., 8.]]];