nn.scan over layers

Scan across depth

Build a stack of N Dense-and-ReLU layers using nn.scan. The activation flows from one layer to the next, and each layer has its own kernel and bias. Unlike an RNN that reuses one cell's weights, this scan selects a different parameter slice at every step.

x_0 = x
x_(k+1) = relu(x_k @ kernel[k] + bias[k]), for k = 0, ..., N-1
return x_N

Each activation has shape (F,). There is no residual addition, normalization or Transformer block in this exercise.

Store a layer axis in each parameter array

Define a compact Block(features=F) that calls one unnamed nn.Dense(F), applies nn.relu and returns (new_x, None). Transform it with:

ScanBlock = nn.scan(
    Block,
    variable_axes={"params": 0},
    split_rngs={"params": True},
    length=N,
)

variable_axes={"params": 0} adds a leading layer axis to every parameter array. With the default Dense bias, the initialized tree has:

params
  Dense_0
    kernel  (N, F, F)
    bias    (N, F)

Layer k reads kernel[k] and bias[k]. There are two arrays containing N * (F * F + F) scalar parameters. The layer axis is a storage layout; it does not make the layers execute in parallel, because each layer consumes the previous layer's output.

split_rngs={"params": True} gives each layer its own initialization key. Random kernels can therefore start differently. Biases still start at zero because their initializer ignores the random key.

If the stream were left unsplit, the same-shaped layers would start with equal values, but they would still occupy separate parameter slices. Equal initialization is not weight tying, and those slices can receive different gradients. To share one parameter set across steps, use a broadcast collection instead of a scanned parameter axis.

Carry the activation, not an output history

x is the carry. The block returns the new activation as its first result and None as its second, so no per-layer history is collected. The scanned call returns (final_activation, None).

Carry structure, shape and dtype must stay fixed. Therefore x must already have length F; this is a scan-shape constraint, not a residual connection. The extra input is None, so pass length=N explicitly rather than relying on an input time axis to set the loop length.

A two-layer example

With known weights, the same computation is easy to check directly:

import jax.numpy as jnp

x = jnp.array([1., -2.])
kernels = jnp.array([
    [[1., 2.], [0., 1.]],
    [[1., -1.], [2., 1.]],
])
biases = jnp.array([[0., 1.], [-1., 0.]])
first = jnp.maximum(x @ kernels[0] + biases[0], 0)
second = jnp.maximum(first @ kernels[1] + biases[1], 0)

assert jnp.array_equal(first, jnp.array([1., 1.]))
assert jnp.array_equal(second, jnp.array([2., 0.]))
assert kernels.shape == (2, 2, 2)
assert biases.shape == (2, 2)

The exercise initializes its weights rather than using these values. The example shows slice order, biases and the activation after each layer.

What scanning changes

A fixed Python loop inside a JAX trace repeats the layer computation in the traced program. Scan represents it as a loop, which can reduce compilation overhead for deep stacks. It does not guarantee constant compilation time or faster execution. Parameter storage and the number of layer evaluations still grow with N.

Both separate layer dictionaries and stacked parameter arrays can be checkpointed. Stacking is a convenient organization for this transform, not a prerequisite for saving model weights.