NNX Scan Layers

Pass one activation through several layers

Start with a vector x of length F. Each layer applies its own linear transformation followed by ReLU, then passes the result to the next layer:

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

nnx.scan expresses this repeated operation as a loop. The activation is the carry: the value passed from one iteration to the next. Unlike scanning a sequence with one shared recurrent cell, this scan consumes a different parameter set at each layer.

Stack parameters along a layer axis

For N width-F linear layers, store kernels as (N, F, F) and biases as (N, F). A mapped constructor builds this representation. Its body describes one layer; split_rngs prepares random-key branches and nnx.vmap maps them through that constructor.

The result is one module-shaped object containing stacked state, not a Python list. Scan slices its array state for each iteration. The object itself is not indexed with blocks[i]. Read parameter arrays with blocks.kernel[...] and blocks.bias[...].

A complete two-layer example

This example assigns known parameters after construction so the result is easy to check. For the exercise itself, retain the default random kernels and zero biases.

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 forward(x, block):
    return jax.nn.relu(block(x))

x = jnp.array([-1., 2.])
# Layer 0: relu([-1, 2] + [0, 1]) = [0, 3].
# Layer 1: relu(2 * [0, 3] + [1, -1]) = [1, 5].
assert jnp.array_equal(forward(x, blocks), jnp.array([1., 5.]))

in_axes=(nnx.Carry, 0) matches the arguments in order: x is carried, and block state is sliced along axis 0. out_axes=nnx.Carry says the body returns the next carry, and the whole scan returns the final one. It does not return all intermediate activations.

To collect activations too, both the body and output specification must change: return (next_x, next_x) and set out_axes=(nnx.Carry, 0). Changing only out_axes does not make the body return a tuple. This exercise requests only the final activation.

What scanning changes

A Python loop inside a JIT-traced function can expand the computation graph with each repeated layer. Scan represents a loop body once and can reduce compilation work for deep stacks. It does not guarantee constant compile time, a particular speedup, or one kernel launch. Measure compilation and execution separately for your workload.

Parameter storage is still N times the storage of one layer: this stack contains N*(F*F + F) scalar parameters. Stacking them does not remove optimizer state associated with each parameter. This task runs a forward pass only and allocates no optimizer.

Layers must have compatible structure and parameter shapes to stack this way. The carry must retain its shape and dtype. A network that changes widths may need separate scans for compatible groups of layers.

Use the specified split-key construction to reproduce initialization. An unsplit RNG may lack the mapped axis; it does not automatically create N identical layers. Deliberately reusing a key can repeat random weights, but applying the same transformation N times still need not equal applying it once. Equal initial values also do not imply shared parameter storage. Default biases are all zero even with split keys.