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 an optimized case when performing the softmax over the last dimension. (#195)

Browse Source
This commit is contained in:
Laurent Mazare
2023-07-18 18:59:50 +02:00
committed by GitHub
parent 6623c227d8
commit fd55fc9592
1 changed files with 24 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -785,6 +785,19 @@ fn divide_by_sum_over_dim<T: WithDType>(s: &mut [T], shape: &Shape, dim: usize)
let elem_per_slice = dims[dim]; let elem_per_slice = dims[dim];
let prod_pre_dim = dims[..dim].iter().product(); let prod_pre_dim = dims[..dim].iter().product();
let prod_post_dim = dims[dim + 1..].iter().product(); let prod_post_dim = dims[dim + 1..].iter().product();
if prod_post_dim == 1 {
for pre_idx in 0..prod_pre_dim {
let mut sum = 0f64;
let idx = pre_idx * elem_per_slice;
for v in s[idx..idx + elem_per_slice].iter() {
sum += v.to_f64();
}
let sum = T::from_f64(sum);
for v in s[idx..idx + elem_per_slice].iter_mut() {
*v /= sum
}
}
} else {
for pre_idx in 0..prod_pre_dim { for pre_idx in 0..prod_pre_dim {
for post_idx in 0..prod_post_dim { for post_idx in 0..prod_post_dim {
let mut sum = 0f64; let mut sum = 0f64;
@ -801,6 +814,7 @@ fn divide_by_sum_over_dim<T: WithDType>(s: &mut [T], shape: &Shape, dim: usize)
} }
} }
} }
}
Ok(()) Ok(())
} }