Skip to main content

vmap

Function vmap 

Source
pub fn vmap<I, O, F, R>(f: F) -> Vmap<I, O, F>
where I: Tree, O: Tree, F: Fn(&I) -> R, R: IntoResult<Ok = O>,
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]);