NNX Composed Transforms

Two axes, two jobs

A batch contains independent input examples. A stack of layers contains successive operations: layer 1 consumes the output of layer 0, and so on. This exercise combines nnx.vmap over examples with nnx.scan over layers in a small feed-forward network.

For N layers of width F, store kernels as (N, F, F) and biases as (N, F). The input batch has shape (B, F). Every sample uses the same N parameter sets; layers have their own parameters.

For each sample, start with h=x and repeat, in layer order:

h = relu(h @ kernel[layer] + bias[layer])

Return only the final activation from each sample. There is no residual connection, normalization, dropout or training step in this exercise.

Scan layers inside a batch map

The inner scan uses in_axes=(nnx.Carry, 0). Its first argument is the activation carried from one layer to the next. Its second argument is the module state, sliced along the leading layer axis. The carry keeps shape (F,) and the same dtype throughout the scan.

out_axes=nnx.Carry returns the final carry, rather than stacking every layer's activation. The outer map uses in_axes=(0, None): take one input row at a time and share the whole layer stack across samples. Its default output axis is 0, giving a result of shape (B, F).

A complete two-layer example

Here we replace random weights with known values so you can check each step. The first layer adds [0, 1] and applies ReLU. The second doubles that activation, adds [1, -1], and applies ReLU again.

import jax
import jax.numpy as jnp
from flax import nnx

@nnx.split_rngs(splits=2)
@nnx.vmap(in_axes=(0,), out_axes=0)
def make_block(rngs):
    return nnx.Linear(2, 2, rngs=rngs)

blocks = make_block(nnx.Rngs(0))
blocks.kernel[...] = jnp.stack([jnp.eye(2), 2 * jnp.eye(2)])
blocks.bias[...] = jnp.array([[0., 1.], [1., -1.]])

@nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)
def scan_layers(x, block):
    return jax.nn.relu(block(x))

@nnx.vmap(in_axes=(0, None))
def batched(x, blocks):
    return scan_layers(x, blocks)

x = jnp.array([[-1., 2.], [3., -2.]])
# After layer 0: [[0, 3], [3, 0]].
out = batched(x, blocks)
assert jnp.allclose(out, jnp.array([[1., 5.], [7., 0.]]))
assert jnp.array_equal(out.reshape(-1), jnp.array([1., 5., 7., 0.]))

Keep initialization and axes consistent

Use the same constructor recipe for the task, with splits=N and nnx.Linear(F, F, rngs=rngs). Keep its default random kernels and zero biases. The outer split_rngs prepares key branches; the constructor's vmap maps those branches and stacks the resulting parameters. A loop using seeds seed+i need not produce the same initialized weights.

Do not confuse the constructor's map over layers with the forward map over samples. Passing (F,) to batched would map its feature axis and send scalars into a computation that expects feature vectors. Pass a two-dimensional (B, F) batch, even when B=1.

For this computation, scanning a batched activation through the layers can give the same result. There is no mutable batch state or randomness in the forward pass. That equivalence should not be assumed for arbitrary stateful modules. Neither ordering guarantees faster execution or a single fused kernel; this task uses scan inside vmap to teach the axes.