use smallvec::smallvec;
use crate::{Element, Recordable, Shape, Tensor};
use super::{Cotangents, Operation, Reads, binary};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct Powf;
impl Powf {
pub(crate) fn arity(&self) -> usize {
2
}
pub(crate) fn reads(&self) -> Reads {
Reads {
operands: [true, true],
output: true,
}
}
pub(crate) fn infer_shape(&self, operands: &[Shape]) -> Shape {
let (base, exponent) = binary(operands);
assert_eq!(base, exponent, "powf requires operands of equal shapes");
base.clone()
}
}
impl Powf {
pub(crate) fn forward<E: Element>(&self, operands: &[&Tensor<E>]) -> Tensor<E> {
let (&base, &exponent) = binary(operands);
base.powf(exponent.clone())
}
}
impl<Rule: Recordable> Operation<Rule> for Powf {
fn backward(&self, operands: &[&Rule], output: &Rule, gradient: &Rule) -> Cotangents<Rule> {
let (&base, &exponent) = binary(operands);
let lowered = exponent.clone() - exponent.one_like();
let base_cotangent = gradient.clone() * exponent.clone() * base.powf(lowered);
let exponent_cotangent = gradient.clone() * output.clone() * base.ln();
smallvec![Some(base_cotangent), Some(exponent_cotangent)]
}
}
#[cfg(test)]
#[path = "tests/powf_tests.rs"]
mod tests;