Skip to main content

vjp_flat

Function vjp_flat 

Source
pub fn vjp_flat(
    f: impl FnOnce(&[Array]) -> Result<Vec<Array>, Error>,
    primals: &[Array],
    cotangents: &[Array],
) -> Result<(Vec<Array>, Vec<Array>), Error>
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]);