Skip to main content

jvp_flat

Function jvp_flat 

Source
pub fn jvp_flat(
    f: impl FnOnce(&[Array]) -> Result<Vec<Array>>,
    primals: &[Array],
    tangents: &[Array],
) -> Result<(Vec<Array>, Vec<Array>)>
Expand description

Forward mode over positional arrays (the flat form of jvp): runs f on primals and returns its outputs together with the Jacobian-vector product, one tangent per output, for one tangent per primal.

use nuvec_core::{Array, Result, jvp_flat};

let x = Array::from(0.0f32);
let (_, tangents) = jvp_flat(|xs| Ok(vec![xs[0].sin()]), &[x], &[Array::from(1.0f32)])?;
assert_eq!(tangents[0].item::<f32>()?, 1.0); // cos(0)