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)