Skip to content

← JAX Vectorization & Control Flow step 7 of 25

Medium Primitives

Scan over Layer Stack

Jump to the task

Carry an activation through a stack of matrices

Repeated blocks can be represented as a stack of layer parameters. A Python loop inside jit unrolls those layer applications during tracing; lax.scan expresses them as a loop with a shared body instead. This can reduce compilation cost for deep stacks. It does not make parameter storage, execution time or gradient memory independent of depth.

This works when the layers have compatible parameter shapes and every layer preserves the carry's shape and dtype. Here each layer is a square matrix and the carry is a vector of fixed length.

Worked mini-example

import jax.numpy as jnp
from jax import lax

weights = jnp.array([
    [[1., 2.], [0., 1.]],
    [[2., 0.], [1., 1.]],
])

def layer(y, w):
    return w @ y, None

x = jnp.array([1., 3.])
final, _ = lax.scan(
    layer, x, weights
)
# final = [14.0, 10.0]

The weight stack has shape (2, 2, 2): two layers, each a 2×2 matrix. Scan consumes one matrix along axis 0 at each step:

Layer Matrix New carry
0 [[1, 2], [0, 1]] [7, 3]
1 [[2, 0], [1, 1]] [14, 10]

The carry is the activation vector, always shape (2,). Returning None as the second result avoids collecting an output at every layer.

The first matrix acts first, so the result is weights[1] @ (weights[0] @ x). Reversing these matrices gives [10, 4], a different answer. Matrix multiplication is generally not commutative. Applying every matrix independently to the original x would also be a different computation, producing one vector per layer.

These layers are linear maps with no bias or activation. Their composition is still linear. This example teaches how scan threads state through a sequence, not how to construct a nonlinear network.

Common pitfalls

  • Returning only the activation. Even when you don't need per-step outputs, you must return (new_carry, None): not just new_carry.
  • Wrong argument order in body. lax.scan calls body(carry, x_i). Here carry is the activation y and x_i is the matrix w. Keep the multiplication order w @ y; y @ w transposes the intended linear map for a one-dimensional y.
  • Confusing xs axis. lax.scan always scans xs along axis 0. Your weight stack must have layers along axis 0 (weights[i] = layer i).

Problem

Implement scan_apply_layers(weights, x) with lax.scan. Start with x as the carry and apply each matrix in order using new_y = w @ y.

The real floating-point input arrays have matching dtypes. The weight stack has shape (N, d, d) and x has shape (d,), with positive d. Return only the final vector of shape (d,). A stack with zero layers returns x unchanged.

Do not add bias, ReLU or another activation. Negative values must remain part of the linear calculation. Use scan rather than a Python loop; output tests verify the result but cannot enforce the loop primitive.