jax4j / NablaTensor

Docs / jax4j-core / com.nablatensor.jax4j.tensor

final class

Vmap

Primitive batching with arbitrary mapped axes. A body is evaluated once, never once per example.

View source ↗Full module reference on javadoc.io ↗

Methods

static Tensor apply(UnaryOperator<Tensor> body, Tensor input)
static UnaryOperator<Tensor> of(UnaryOperator<Tensor> body)
static Tensor apply(Function<List<Tensor>, Tensor> body, List<Tensor> inputs, List<Integer> inAxes, int outAxis)

Null input axes mean shared arguments; negative axes count from the end.

static R applyTree(Function<P, R> body, P inputs, List<Integer> inAxes, List<Integer> outAxes)

Input/output axes are in deterministic pytree leaf order; null output axes require an unmapped result.