NNX vmap over Batch

Apply one model to many inputs

Suppose a function accepts one feature vector of shape (D_in,) and returns (D_out,). Vectorization adds a batch axis: apply that same function to B rows and collect a result shaped (B, D_out).

This exercise uses nnx.vmap with one shared linear layer. Linear layers already accept batched inputs directly, so model(x_batch) would produce the same numbers here. The purpose is to learn the axis specification for an NNX module argument, which also works with custom per-example computations that do not directly accept a batch.

Read the axes one argument at a time

For a function fwd(model, x), use in_axes=(None, 0):

• None leaves the model's parameter arrays unmapped. Every example uses the same kernel and bias.
• Zero maps the input's leading axis. Inside fwd, x has shape (D_in,).

The default output axis is zero, so vmap stacks the output vectors into (B, D_out). It does not create B sets of parameters or update the model once per row as a sequential training loop would.

A complete example

import jax.numpy as jnp
from flax import nnx

model = nnx.Linear(2, 2, rngs=nnx.Rngs(0))
model.kernel[...] = jnp.array([[1.0, 2.0], [3.0, 4.0]])
model.bias[...] = jnp.array([0.5, -0.5])

@nnx.vmap(in_axes=(None, 0))
def fwd(model, x):
    return model(x)

batch = jnp.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
expected = jnp.array([[1.5, 1.5], [3.5, 3.5], [4.5, 5.5]])
assert jnp.allclose(fwd(model, batch), expected)
assert jnp.allclose(fwd(model, batch), model(batch))

Each output row is input_row @ kernel + bias. The bias is added to every row, but it remains a single parameter vector. A batch loss can accumulate gradients from all rows into those same shared parameters.

Avoid confusing a batch with an ensemble

Mapping axis zero of the model would require appropriately stacked parameter state. It describes multiple parameter sets, not one model shared by a batch. (0, None) applies such models to a shared input; (0, 0) pairs them with inputs of matching mapped size.

For this task, do not map the model or construct it inside fwd. Build it once, then pass it as the unmapped first argument.

A one-dimensional input still has axis zero, but mapping it supplies scalars to the function rather than feature vectors. For one example, use shape (1, D_in), for example x[None, :]. Do not confuse a missing batch dimension with a nonexistent axis.

State and tracing

nnx.vmap handles NNX module state as part of the transformation. This linear forward pass does not mutate state. For stateful layers, deciding which state is shared and which is per-example needs additional care; vectorization does not imply sequential updates to shared state.

The mapped body is traced. Data-dependent Python branches do not become valid just because vmap adds a batch axis; use suitable JAX control flow. Vectorization also does not by itself guarantee a speedup for a workload.

The diagram expands the model into separate kernel and bias arguments to show their array shapes. Both are shared. For eight input rows of width three and output width four, the output is (8, 4), flattened to 32 values. In the code, NNX accepts the whole model as one argument.