pub fn vmap<F>(
f: F,
in_axes: &[Option<isize>],
out_axes: &[isize],
) -> impl Fn(&[Array]) -> Result<Vec<Array>>Expand description
Returns f vectorized over a batch axis.
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::{Array, Result, matmul, vmap};
// A matrix-vector product per example, sharing the matrix.
let mv = vmap(|xs| Ok(vec![matmul(&xs[0], &xs[1])?]), &[None, Some(0)], &[]);
let m = Array::from_slice(&[1.0f32, 0.0, 0.0, 2.0], &[2, 2])?;
let vs = Array::from_slice(&[1.0f32, 1.0, 3.0, 4.0], &[2, 2])?; // two vectors
assert_eq!(mv(&[m, vs])?[0].to_vec::<f32>()?, [1.0, 2.0, 3.0, 8.0]);