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

Add the affine transformation.

Browse Source
This commit is contained in:
laurent
2023-06-20 21:51:35 +01:00
parent e04d0ac1ac
commit f1f372b13e
3 changed files with 71 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
src/storage.rs
View File

@ -172,6 +172,32 @@ impl Storage {
}
}
pub(crate) fn affine_impl(
&self,
shape: &Shape,
stride: &[usize],
mul: f64,
add: f64,
) -> Result<Self> {
// TODO: Different code path for the contiguous case?
match self {
Storage::Cpu(storage) => match storage {
CpuStorage::F32(storage) => {
let index = StridedIndex::new(shape.dims(), stride);
let mul = mul as f32;
let add = add as f32;
let data = index.map(|i| storage[i] * mul + add).collect();
Ok(Storage::Cpu(CpuStorage::F32(data)))
}
CpuStorage::F64(storage) => {
let index = StridedIndex::new(shape.dims(), stride);
let data = index.map(|i| storage[i] * mul + add).collect();
Ok(Storage::Cpu(CpuStorage::F64(data)))
}
},
}
}
fn unary_impl<B: UnaryOp>(&self, shape: &Shape, stride: &[usize]) -> Result<Self> {
// TODO: Different code path for the contiguous case?
match self {