Multiple PRNG Streams (params, dropout)

Flax Linen uses named random-number streams to separate tasks such as weight initialization and dropout. This exercise uses one key for Dense parameters and another for a dropout mask, then computes a single forward pass.

Keys are inputs, not stored masks

A stream is a name mapped to a JAX random key. During a call, Linen derives keys using that stream, the module path and a call counter. You supply the stream keys at the boundary:

variables = model.init(
    {"params": init_key, "dropout": drop_key}, x, deterministic=False
)
output = model.apply(
    variables, x, deterministic=False, rngs={"dropout": drop_key}
)

params supplies the keys used to initialize weights. Dropout requests the dropout stream when it samples. The returned variables for this model contain only the params collection: an RNG stream is not a collection of stored variables, and dropout does not save its mask there.

model.init(key, ...) is shorthand for supplying a params stream. If a requested stream is absent but params is available, Linen can fall back to that key. Thus active dropout does not always make single-key initialization fail. Explicit separate streams make the intended sources of randomness clear, and this exercise requires them.

The model and its output

class DropoutModel(nn.Module):
    drop_rate: float

    @nn.compact
    def __call__(self, x, deterministic):
        x = nn.Dense(features=x.shape[-1])(x)
        return nn.Dropout(
            rate=self.drop_rate, deterministic=deterministic
        )(x)

Let z be the Dense output and r the dropout rate. For 0 < r < 1 and deterministic=False, each element is independently kept with probability 1 - r. A kept element becomes z[i] / (1 - r); a dropped element becomes zero. The expected output equals z, but the number dropped varies by mask.

deterministic=True returns z unchanged, regardless of rate. With dropout enabled, rate zero also returns z, while rate one returns all zeros. These no-sampling cases need no dropout key. Sampling with no available dropout or fallback params stream raises a missing-RNG error.

Reproducibility and fresh randomness

The same key, shapes and call structure reproduce the same mask. This is useful for tests and is intentional in this exercise. In a training loop, split a carried key to obtain a new dropout key for each step:

rng, drop_key = jax.random.split(rng)
output = model.apply(
    variables, batch, deterministic=False, rngs={"dropout": drop_key}
)

Distinct keys do not guarantee distinct masks for a small array. Also, Linen's internally derived mask need not match a direct Bernoulli draw using the top-level key. Use the specified module and naming for these tests.