pub fn vjp(
f: impl FnOnce(&[Array]) -> Result<Vec<Array>>,
primals: &[Array],
cotangents: &[Array],
) -> Result<(Vec<Array>, Vec<Array>)>Expand description
Reverse mode: 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::{Array, Result, multiply, vjp};
let x = Array::from_slice(&[1.0f32, 2.0, 3.0], &[3])?;
let ones = Array::from_slice(&[1.0f32; 3], &[3])?;
let (_, grads) = vjp(|xs| Ok(vec![multiply(&xs[0], &xs[0])?]), &[x], &[ones])?;
assert_eq!(grads[0].to_vec::<f32>()?, [2.0, 4.0, 6.0]);