NNX Batched Init

Initialize a model axis

Batched initialization builds several parameter sets with the same structure and stores them along a leading model axis. For N linear models of width F, the kernel has shape (N, F, F) and the bias has shape (N, F).

The result is one NNX module-shaped object containing stacked arrays, rather than a Python list of N model objects. This representation is useful when later transformations map over the model axis. It still stores N sets of parameters; vectorization does not remove that cost.

Split keys, then map the constructor

The constructor body describes one nnx.Linear(F, F, rngs=rngs). Two decorators give it a model axis:

• The outer nnx.split_rngs(splits=N) prepares RNG state with a leading axis of N key branches.
• The inner nnx.vmap(in_axes=(0,), out_axes=0) maps that RNG axis into the single-model constructor and stacks the resulting state.

Each lane receives its own random-key branch. The kernel initializer uses those branches for separate draws, while the default biases are all zero. Separate draws do not mean every parameter must differ. Repeating the same key deliberately would reproduce the same draw.

A loop using seeds seed+i or sequential keys follows a different random construction and need not yield the same weights. For this exercise, use the specified decorators and a single nnx.Rngs(int(seed)).

A complete example

import jax.numpy as jnp
from flax import nnx

@nnx.split_rngs(splits=3)
@nnx.vmap(in_axes=(0,), out_axes=0)
def make_model(rngs):
    return nnx.Linear(2, 2, rngs=rngs)

ensemble = make_model(nnx.Rngs(7))
assert ensemble.kernel[...].shape == (3, 2, 2)
assert ensemble.bias[...].shape == (3, 2)
assert jnp.all(ensemble.bias[...] == 0)
summary = jnp.array([3, 2, ensemble.kernel[...].shape[0]])
assert jnp.array_equal(summary, jnp.array([3, 2, 3]))

The three-by-two-by-two kernel has 12 entries, and the bias has six. The three-element summary reports dimensions; it is not a parameter count or a collection of predictions. It stays the same for any seed when N and F are unchanged, even though the initialized kernels change.

Read wrapped parameter arrays with parameter[...]. The kernel's leading dimension comes from the mapped constructor, not a batch of input data. No input examples are supplied during this task.

Keep the ensemble intact

To apply these models to one shared input later, use a mapped forward with in_axes=(0, None): model state is mapped, input is shared. Averaging the parameter axis would create one different parameter set and discard the individual models. It is not a general substitute for evaluating models separately and combining their predictions.

nnx.vmap handles the module structure and array state for this recipe. An unsplit RNG lacks the mapped key axis it expects. Use NNX-aware transformations when composing operations that take these modules. Actual execution speed depends on compilation, hardware and model size.