NNX vmap over Ensemble

Many models, one input

An ensemble has several parameter sets. Here each member is a linear layer with the same input and output widths, initialized from its own random-key branch. Apply all members to one shared input vector and keep their outputs separate.

This is a forward-pass exercise. It neither trains the models nor averages predictions, and the outputs do not establish calibrated uncertainty. Vectorization can help execute the computation, but parameter storage and arithmetic still grow with the number of models.

Put the model index on axis zero

For N models mapping D input features to F output features, the stacked kernel has shape (N, D, F) and the stacked bias has shape (N, F). NNX represents these arrays inside one module-shaped object.

A forward function with in_axes=(0, None) maps the model state along axis zero and leaves the input vector unmapped. Inside the function, each model sees a (D, F) kernel, an (F,) bias and the same (D,) x. The output is (N, F), with one row per model.

Compare these patterns for a function fwd(model, x):

Pairing models and inputs does not by itself implement mixture-of-experts routing. Also, do not simply call the stacked object as an ordinary Linear layer: use the mapped forward to give each parameter slice its intended model meaning.

Initialize the stacked parameters

Use two decorators on a single-model constructor. The outer nnx.split_rngs(splits=N) prepares a mapped RNG axis. The inner nnx.vmap(in_axes=(0,), out_axes=0) constructs a model in each lane and stacks its parameter state on axis zero.

The order matters: split the RNG state before mapping the constructor. Split keys provide separate random draws; repeatedly using the same key would reproduce the same initialization. Not every parameter becomes different: the default linear biases are all zero. Distinct random draws also do not guarantee that every output differs, for example on zero input.

Use the exact seeded construction below. A loop with sequential RNG draws or seeds seed+n can produce valid models, but it need not reproduce these particular initialized weights.

A complete example

This creates two models and then sets their parameters to known values so the shared-input mapping is easy to check:

import jax.numpy as jnp
from flax import nnx

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

ensemble = make_model(nnx.Rngs(0))
ensemble.kernel[...] = jnp.array([[[1.0], [2.0]], [[3.0], [4.0]]])
ensemble.bias[...] = jnp.array([[0.5], [-0.5]])

@nnx.vmap(in_axes=(0, None))
def fwd(model, x):
    return model(x)

x = jnp.array([1.0, 2.0])
assert jnp.allclose(fwd(ensemble, x), jnp.array([[5.5], [10.5]]))

The first member computes 1*1 + 2*2 + 0.5; the second computes 1*3 + 2*4 - 0.5. The result contains both values, not their mean. Altering one member's parameter slice changes only that member's output.

The diagram expands the module into kernel and bias arrays. Both map their leading model axis, while x is shared. Five models with output width four produce (5, 4), flattened to 20 values for this task.