pub fn vjp_flat(
f: impl FnOnce(&[Array]) -> Result<Vec<Array>>,
primals: &[Array],
cotangents: &[Array],
) -> Result<(Vec<Array>, Vec<Array>)>Expand description
Reverse mode over positional arrays (the flat form of vjp): runs f on primals
and returns its outputs together with the vector-Jacobian product, one gradient per
primal, for one cotangent per output.
The results are lazy arrays. A primal that the outputs do not depend on gets a zero gradient.
use nuvec_core::{Result, array, vjp_flat};
let x = array![1.0, 2.0, 3.0];
let ones = x.ones_like();
let (_, grads) = vjp_flat(|xs| Ok(vec![xs[0].multiply(&xs[0])]), &[x], &[ones])?;
assert_eq!(grads[0].to_vec::<f32>()?, [2.0, 4.0, 6.0]);