Composed lifts: nn.scan + nn.vmap

Batch a stack of layers

Combine two Linen transformations: scan across N layers, then map that stack across B input examples. Each layer owns a kernel and bias, and every example uses the same complete stack.

The input and final output both have shape (B, F). For each row b, compute:

h = x_batch[b]
for k = 0, ..., N-1:
    h = relu(h @ kernel[k] + bias[k])
output[b] = h

This is a Dense/ReLU stack. It has no residual additions, attention, normalization or interaction between different examples.

Build the transformations from the inside out

Define a compact Block(features=F) that applies one unnamed nn.Dense(F) followed by nn.relu, returning (new_x, None). First create a per-example stack:

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

The scan carries one activation of shape (F,) from layer to layer. variable_axes places a layer axis on the parameters; split keys let random kernels initialize differently. Dense biases still initialize to zero. No history is collected because the block's second result is None.

Then map the entire stack across input rows:

BatchedScanLayer = nn.vmap(
    ScanLayer,
    in_axes=0,
    out_axes=0,
    variable_axes={"params": None},
    split_rngs={"params": False},
)

The outer transform maps the activation's batch axis while keeping the whole parameter collection shared. Its unsplit parameter RNG matches that sharing policy. The extra None input has no array leaves to map.

Inspect the parameter layout

Both transformations configure the same params collection at different levels. They do not configure different collections. The resulting tree is:

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

There is no B axis on either array. The scalar parameter count is N * (F * F + F), independent of batch size. The returned pair is (final, None), with final shaped (B, F).

A mapped parameter axis on the outer transform would instead give each example a separate stack. Splitting an RNG alone does not request that storage layout; parameter axes and random streams are separate settings.

Composition order and sharing are different choices

Use the prescribed outer vmap and inner scan for this exercise. Reversing the transforms does not automatically create per-example weights: a scan of a shared-parameter batched block can also describe this computation. The variable-axis settings determine sharing.

Different compositions can change axis interpretation, parameter layout or random-key routing. Do not assume arbitrary transformations commute or that changing their order preserves seeded initialization. Matching mathematical operations with the same weights is a separate question from matching how those weights are initialized.

A two-example, two-layer calculation

This explicit calculation applies the same two parameter slices to both examples:

import jax.numpy as jnp

x = jnp.array([[1., -2.], [0., 1.]])
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)
final = jnp.maximum(first @ kernels[1] + biases[1], 0)

assert jnp.array_equal(first, jnp.array([[1., 1.], [0., 2.]]))
assert jnp.array_equal(final, jnp.array([[2., 0.], [3., 2.]]))
assert jnp.array_equal(final.reshape(-1), jnp.array([2., 0., 3., 2.]))

This uses known weights to show the axes. The exercise initializes its own weights from the supplied seed.