Skip to main content

vjp

nuvec_core

Function vjp 

Source 
pub fn vjp<P, O, F, R>(f: F, primals: &P, cotangents: &O) -> Result<(O, P)>
where
    P: Tree,
    O: Tree,
    F: FnOnce(&P) -> R,
    R: IntoResult<Ok = O>,
Expand description

Reverse mode over trees: f’s output at primals and the vector-Jacobian product for cotangents (a tree like the output: same paths, shapes and dtypes, else an error naming the path), as a tree like primals (zeros for leaves that are not differentiated).

use nuvec_core::{Array, Result, array, vjp};

let x = array![1.0, 2.0, 3.0];
let (_, g) = vjp(|x: &Array| x.multiply(x), &x, &x.ones_like())?;
assert_eq!(g.to_vec::<f32>()?, [2.0, 4.0, 6.0]);