Struct Vmap 

pub struct Vmap<I, O, F> { /* private fields */ }

A function vectorized over batch axes (vmap): configure it with Vmap::in_axes/Vmap::out_axes, then call it with Vmap::call or turn it into a closure with Vmap::into_fn (to compose with compile and the other transforms).

Implementations

impl<I: Tree, O: Tree, F> Vmap<I, O, F>

pub fn in_axes(self, axes: impl InAxes) -> Self

The batch axis of each top-level argument (None: shared by every example), or one axis for all: in_axes((None, 0)), in_axes(1). Negative axes count from the end.

pub fn out_axes(self, axes: impl OutAxes) -> Self

The axis where each top-level output gets the batch dimension, or one for all.

pub fn call<R>(&self, inputs: &I) -> Result<O>

where F: Fn(&I) -> R, R: IntoResult<Ok = O>,

f batched over inputs.

pub fn into_fn<R>(self) -> impl Fn(&I) -> Result<O>

where F: Fn(&I) -> R, R: IntoResult<Ok = O>,

The configured batched function as a closure, which composes with compile, grad and the other transforms.