pub fn vmap_flat<F>(
f: F,
in_axes: &[Option<i32>],
out_axes: &[i32],
) -> impl Fn(&[Array]) -> Result<Vec<Array>>Expand description
Returns f vectorized over a batch axis, over positional arrays (the flat form of
vmap).
in_axes gives each argument’s batch axis (None for arguments shared by every
example) and out_axes each output’s; empty slices mean axis 0 everywhere. Negative
axes count from the end. Outputs that do not depend on a batched argument are broadcast
along the batch axis. Arrays cannot be evaluated inside f.
use nuvec_core::{Result, array, vmap_flat};
// A matrix-vector product per example, sharing the matrix.
let mv = vmap_flat(|xs| Ok(vec![xs[0].matmul(&xs[1])]), &[None, Some(0)], &[]);
let m = array![[1.0, 0.0], [0.0, 2.0]];
let vs = array![[1.0, 1.0], [3.0, 4.0]]; // two vectors
assert_eq!(mv(&[m, vs])?[0].to_vec::<f32>()?, [1.0, 2.0, 3.0, 8.0]);