pub fn vmap<I, O, F, R>(f: F) -> Vmap<I, O, F>Expand description
Returns f vectorized over a batch axis of its inputs (axis 0 of every top-level
argument and output by default; see Vmap::in_axes and Vmap::out_axes).
f takes a tree (an array, a tuple of arguments, a struct) and returns one. Outputs
that do not depend on a batched input are broadcast along the batch axis. Arrays cannot
be evaluated inside f.
use nuvec_core::{Array, Result, array, vmap};
// A matrix-vector product per example, sharing the matrix.
let mv = vmap(|(m, v): &(Array, Array)| m.matmul(v)).in_axes((None, 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.call(&(m, vs))?.to_vec::<f32>()?, [1.0, 2.0, 3.0, 8.0]);