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

Add binary and ternary custom ops. (#217)

Browse Source
This commit is contained in:
Laurent Mazare
2023-07-21 18:29:50 +02:00
committed by GitHub
parent 4a100875bf
commit 5cc843550d
5 changed files with 209 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
candle-core/tests/custom_op_tests.rs
View File

@ -94,10 +94,10 @@ impl CustomOp1 for EluWithBackward {
self.0.cpu_fwd(s, l)
}
fn bwd(&self, arg: &Tensor, _res: &Tensor, grad_res: &Tensor) -> Result<Tensor> {
fn bwd(&self, arg: &Tensor, _res: &Tensor, grad_res: &Tensor) -> Result<Option<Tensor>> {
let alpha = self.0.alpha;
let bwd = arg.custom_op1(EluBackward { alpha })?;
grad_res.mul(&bwd)
Ok(Some(grad_res.mul(&bwd)?))
}
}