Skip to content

← Flax step 97 of 100

Medium Primitives

Multiple Mutable Collections

Jump to the task

Combine a Dense layer, dropout and a persistent call counter in one Flax Linen module. The exercise shows how stored variables and random keys enter an apply call through different arguments.

Two variable collections, two RNG streams

This model has two stored collections:

  • params contains the Dense kernel and bias.
  • buffers contains a scalar named running_count.

It also uses two RNG streams: params for initializing Dense weights and dropout for sampling masks. RNG streams and variable collections are separate concepts. There is no stored dropout collection in this model, and creating a collection does not automatically give it an RNG stream.

Declare and update the counter

Inside the compact method, declare:

running_count = self.variable(
    "buffers", "running_count", lambda: jnp.array(0.0)
)

The function in the third argument creates the initial value. Read the array through .value and increment it after the forward calculation:

if not self.is_initializing():
    running_count.value = running_count.value + 1.0

This guard leaves the initialized count at zero. Initialization can write real variable values: without the guard, the returned initial count would be one. Such a write does not disappear or inherently cause an error.

Permit writes and retain the result

Call the model with mutable=["buffers"]:

out, updates = model.apply(
    variables, x, rngs={"dropout": dropout_rng}, mutable=["buffers"]
)

The counter in updates["buffers"] is now one. The original variables mapping still has count zero. To continue counting, carry the updated collection into the next call:

variables = {**variables, "buffers": updates["buffers"]}

Applying with that new mapping increments the count to two. Reusing the original mapping would produce one again. Multiple collections can be writable in one call, but this exercise writes only buffers; params remain read-only. No optimizer update is performed.

Dropout still runs during initialization

The forward path is an unnamed nn.Dense(4) followed by an unnamed nn.Dropout(rate=self.drop_rate, deterministic=False). For an intermediate rate, dropout samples and masks the Dense output during initialization too. It does not wait for a later training call. The wrapper uses only the initialized variables from that first pass.

Split the seed key once, assigning the first result to parameter initialization and the second to dropout. Reuse that dropout key for this exercise's application so the result is reproducible. In a training loop, supply fresh keys when you want fresh randomness. Rate zero passes values through; rate one returns zeros. The counter increments in either case.

Problem

Implement MultiColModel(nn.Module) with a drop_rate field and a compact __call__. Declare the counter, apply Dense then Dropout as above, increment the counter only outside initialization, and return the (N, 4) output.

Implement multi_collections_step(seed, x, drop_rate):

  1. Set rng = jax.random.PRNGKey(int(seed)), then init_rng, dropout_rng = jax.random.split(rng).
  2. Construct MultiColModel(drop_rate=float(drop_rate)). Initialize with {"params": init_rng, "dropout": dropout_rng} and x.
  3. Apply once using the full variables mapping, rngs={"dropout": dropout_rng} and mutable=["buffers"].
  4. Return a JAX array of shape (2,) containing the updated counter and the sum of all output elements, in that order. The counter is 1.0.

x is a finite float32 array of shape (N, in_dim), with both dimensions positive. drop_rate is in [0, 1]. seed is a nonnegative integer-valued scalar that may arrive as a float. Leave x unchanged. Keep the layer ordering and default Dense settings, including bias, because names and key derivation affect the expected seeded output.