Skip to content

← Flax step 81 of 100

Medium Primitives

nn.vmap with shared params

Jump to the task

Apply the same layer to every row

Use nn.vmap to apply a Dense layer independently to each example in a batch. Every row uses the same kernel and bias:

y[b] = x_batch[b] @ kernel + bias

For input shape (B, D) and output width F, the kernel has shape (D, F), the bias has shape (F,), and the output has shape (B, F). There is no activation function, reduction across examples or training step. Return all output rows flattened.

Dense already supports leading batch dimensions. We use a lifted map here to learn how the transform controls parameters separately from input and output axes.

Specify the data axes and variable axes separately

Construct the transformed module class with:

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

in_axes=0 maps over input rows. Each mapped call receives (D,). out_axes=0 places the example axis first in the result, giving (B, F).

variable_axes={"params": None} keeps the parameter arrays unbatched. Both kernel and bias are shared, so the parameter count is D * F + F, independent of the number of examples.

split_rngs={"params": False} leaves the parameter-initialization stream unsplit along the mapped axis. This matches the shared variable collection. RNG handling and parameter storage are separate choices; changing the RNG option alone does not request per-example parameters.

Compare shared and per-example weights

With variable_axes={"params": 0}, the parameter arrays instead have shapes (B, D, F) and (B, F). Each example uses its own slice. Splitting the initialization stream gives the random kernels different keys; leaving it unsplit gives equal initial values in separate slices.

That setup can help build ensembles, but this particular mapping pairs one parameter slice with one input row. It does not automatically apply every model to every example or combine their predictions. This exercise uses shared weights throughout.

Check the computation with known weights

The following example uses explicit parameters to isolate the batching behavior from random initialization:

import jax
import jax.numpy as jnp
import flax.linen as nn

VmappedDense = nn.vmap(
    nn.Dense, in_axes=0, out_axes=0,
    variable_axes={"params": None}, split_rngs={"params": False},
)
model = VmappedDense(features=2)
x = jnp.array([[1., 2.], [-1., 0.], [1., 2.]])
variables = {"params": {
    "kernel": jnp.array([[1., -1.], [2., .5]]),
    "bias": jnp.array([.5, 1.]),
}}
y = model.apply(variables, x)
expected = jnp.array([[5.5, 1.], [-.5, 2.], [5.5, 1.]])
assert jnp.allclose(y, expected)
assert jnp.array_equal(y[0], y[2])

dense = nn.Dense(features=2)
mapped_apply = jax.vmap(lambda row: dense.apply(variables, row))
assert jnp.allclose(mapped_apply(x), y)
assert jnp.allclose(dense.apply(variables, x), y)

Ordinary jax.vmap over a pure apply call does not secretly create parameters. Here it closes over the same initialized variables for every row. Lifted nn.vmap adds control over Linen initialization, variable collections and random streams. Collections and streams other than params would need their own appropriate configuration.

Problem

Implement nn_vmap_forward(seed, x_batch, features):

  1. Convert features to Python integer F. Construct VmappedDense with the exact options above and instantiate it with features=F.
  2. Initialize that module directly with jax.random.PRNGKey(int(seed)) and x_batch. Keep Dense's default kernel initializer and zero bias initializer.
  3. Apply those variables to x_batch and return .reshape(-1).

x_batch is a two-dimensional float array of shape (B, D); B, D and F are positive. Return length B * F, with each row's features together. Keep seed and features static under JIT.

Initialize the transformed Dense directly, without a wrapper module or explicit name, so the root parameter keys remain kernel and bias. Preserve this scope and the initializer choices for the seeded tests.