use smallvec::smallvec;
use crate::{Element, Recordable, Shape, Tensor};
use super::{Cotangents, Operation, Reads, binary};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct Div;
impl Div {
pub(crate) fn arity(&self) -> usize {
2
}
pub(crate) fn reads(&self) -> Reads {
Reads {
operands: [false, true],
output: true,
}
}
pub(crate) fn infer_shape(&self, operands: &[Shape]) -> Shape {
let (left, right) = binary(operands);
assert_eq!(left, right, "division requires operands of equal shapes");
left.clone()
}
}
impl Div {
pub(crate) fn forward<E: Element>(&self, operands: &[&Tensor<E>]) -> Tensor<E> {
let (&left, &right) = binary(operands);
left.clone() / right.clone()
}
}
impl<Rule: Recordable> Operation<Rule> for Div {
fn backward(&self, operands: &[&Rule], output: &Rule, gradient: &Rule) -> Cotangents<Rule> {
let (_, &right) = binary(operands);
smallvec![
Some(gradient.clone() / right.clone()),
Some(-(gradient.clone() * output.clone() / right.clone())),
]
}
}