Skip to content

← Flax step 86 of 100

Medium Primitives

Batched init via jax.vmap

Jump to the task

To compare several random initializations, create several sets of parameters for the same model. In Flax Linen, jax.vmap can map model.init over a batch of random keys while using the same example input for every initialization. The result is one variables tree with a new leading axis on each array.

Map keys, share the input

For a Dense layer with F output features and an input of length four:

model = nn.Dense(features=F)
keys = jax.random.split(jax.random.PRNGKey(seed), K)
x = jnp.ones((4,))
stacked = jax.vmap(model.init, in_axes=(0, None))(keys, x)

in_axes=(0, None) describes the two arguments to init:

  • 0 maps over the keys. Each initialization receives one key.
  • None leaves x unmapped. Each initialization receives the whole vector.

The variables tree has these shapes:

Array Shape
stacked["params"]["kernel"] (K, 4, F)
stacked["params"]["bias"] (K, F)

The new axis indexes models, not input examples. To select model i, slice every leaf with jax.tree_util.tree_map(lambda a: a[i], stacked). The result is a variables mapping that can be passed to model.apply.

Split the key rather than repeat it. Repeating one key repeats the same initialized parameters. Different split keys provide separate random streams for the kernels; the default biases are still all zeros.

With the legacy PRNGKey used here, keys has shape (K, 2). Typed keys created with jax.random.key have a different representation. The important part for vmap is the leading model axis.

What vectorization changes

A Python loop with a fixed iteration count can be traced by JAX. vmap expresses the batch directly instead of unrolling that loop. It does not promise one device per model or a particular speedup.

Here jax.vmap acts on the external model.init function. Linen also has nn.vmap, which manages variables and random streams when transforming a module. This exercise uses the external function and needs only jax.vmap.

Problem

Implement batched_init_param_count(seed, num_inits, features). The inputs are integer-valued scalars, possibly supplied as floats. Cast them to int. seed is nonnegative, and num_inits=K and features=F are positive.

Create nn.Dense(features=F) with its default settings. Split jax.random.PRNGKey(seed) into K keys and initialize the layer with jax.vmap(model.init, in_axes=(0, None)), using jnp.ones((4,)) as the shared input. Read the leading dimension of the resulting kernel array.

Return a floating-point JAX array of shape (3,) containing [num_inits, features, kernel_leading_axis]. For K=4 and F=8, the kernel has shape (4, 4, 8) and the result is [4.0, 8.0, 4.0].

This is a shape-inspection exercise, despite the function's historical name: the return value does not count scalar parameters. Perform the initialization and read its shape rather than hard-coding the third value.