Add the powf op. (#664)

* Add the powf op.

* Cuda kernels and backprop.

* Add a test.
This commit is contained in:
Laurent Mazare
2023-08-29 20:48:18 +01:00
committed by GitHub
parent 2d3fcad267
commit 59b731de99
10 changed files with 103 additions and 0 deletions

View File

@ -105,6 +105,7 @@ impl Tensor {
| Op::Narrow(node, _, _, _)
| Op::Unary(node, _)
| Op::Elu(node, _)
| Op::Powf(node, _)
| Op::CustomOp1(node, _) => {
let (tg, nodes) = walk(node, nodes, already_seen);
track_grad |= tg;
@ -437,6 +438,11 @@ impl Tensor {
*sum_grad = sum_grad.add(&(&grad * relu_grad)?)?
}
Op::Elu(..) => Err(Error::BackwardNotSupported { op: "elu" })?,
Op::Powf(arg, e) => {
let arg_grad = (&(grad * arg.powf(e - 1.)?)? * *e)?;
let sum_grad = grads.or_insert(arg)?;
*sum_grad = sum_grad.add(&arg_grad)?
}
Op::CustomOp1(arg, c) => {
if let Some(arg_grad) = c.bwd(arg, node, &grad)? {
let sum_grad = grads.or_insert(arg)?;