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

Skeleton implementation for softmax.

Browse Source
This commit is contained in:
laurent
2023-06-23 22:00:13 +01:00
parent 5d44e76e3f
commit 8443963d4f
5 changed files with 32 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
src/storage.rs
View File

@ -72,6 +72,13 @@ impl Storage {
}
}
pub(crate) fn divide_by_sum_over_dim(&mut self, shape: &Shape, dim: usize) {
match self {
Storage::Cpu(storage) => storage.divide_by_sum_over_dim(shape, dim),
Self::Cuda(storage) => storage.divide_by_sum_over_dim(shape, dim),
}
}
pub(crate) fn to_dtype(&self, shape: &Shape, stride: &[usize], dtype: DType) -> Result<Self> {
match self {
Storage::Cpu(storage) => {