Skip to content

← Flax NNX step 28 of 100

Medium Primitives

NNX Implement Dropout

Jump to the task

Randomly drop values during training

Dropout replaces selected activations with zero during training. In this exercise, each element has drop probability rate and keep probability q = 1 - rate. Kept values are divided by q; evaluation returns the input unchanged. This convention is called inverted dropout.

For a fixed input value x, the output is x / q with probability q and zero otherwise. Its expected value is therefore q * (x / q) = x. That is an expectation over random masks, not a promise that each sampled output has the same sum or mean as the input. It also does not make a whole nonlinear network's training and evaluation outputs identical.

For example, at rate 0.5, an activation of 3 becomes either 0 or 6. Averaging those two equally likely outcomes gives 3. Evaluation returns 3 directly, without another scale factor.

A complete implementation and example

import jax
import jax.numpy as jnp
from flax import nnx

class MyDropout(nnx.Module):
    def __init__(self, rate, rngs):
        self.rate = rate
        self.rngs = rngs

    def __call__(self, x, deterministic):
        if deterministic or self.rate <= 0.0:
            return x
        keep_prob = 1.0 - self.rate
        key = self.rngs.dropout()
        mask = jax.random.bernoulli(key, p=keep_prob, shape=x.shape)
        return jnp.where(mask, x / keep_prob, 0.0)

model = MyDropout(0.5, nnx.Rngs(0))
x = jnp.ones(8, dtype=jnp.float32)
first = model(x, deterministic=False)
assert jnp.all((first == 0) | (first == 2))
assert int(model.rngs.default.count[...]) == 1
assert jnp.array_equal(model(x, deterministic=True), x)
assert int(model.rngs.default.count[...]) == 1
second = model(x, deterministic=False)
assert int(model.rngs.default.count[...]) == 2

A training call with positive rate consumes one key. Evaluation and a zero-rate call consume none. Two draws can produce the same mask by chance, especially for small inputs; a new key does not guarantee visibly different outputs. Returning x directly in evaluation is sufficient because JAX arrays are immutable.

What the RNG container stores

nnx.Rngs(0) creates a default stream. When no named dropout stream exists, rngs.dropout() falls back to that default stream. It does not create a new named stream on its first call. The stream derives a key by folding its current counter into its root key, then increments the counter. A params call using the same fallback also advances that counter.

Named streams such as nnx.Rngs(params=0, dropout=42) have separate counters. That can separate parameter initialization from dropout draws. For this exercise, use the specified positional seed and its default stream to reproduce the expected samples.

The layer has no trainable nnx.Param, but it does have state: its RNG container holds key and counter variables. Keep the same module to advance that state over multiple training calls. Use NNX-aware transforms such as nnx.jit for a persistent module with RNG updates.

Problem

Implement dropout_forward(seed, x, rate, deterministic):

  1. Define MyDropout as above, storing the rate and RNG container.
  2. Convert the mode input using bool(deterministic >= 0.5).
  3. Build a fresh module with float(rate) and nnx.Rngs(int(seed)).
  4. Return model(x, deterministic=det), preserving the input shape.

The fixtures use nonempty float32 vectors. The elementwise mask also supports higher-rank arrays; it has exactly x.shape, without sharing a mask across any axis. Inputs are finite, and retained values must remain representable after scaling. rate is finite and lies in [0, 1); rate 1 is outside the contract because the scaling would divide by zero. The mode is 0.0 for training or 1.0 for evaluation. seed is nonnegative and integer-valued, possibly supplied as a float.

Each entrypoint call starts a new RNG container, so repeating it with the same inputs returns the same sample. This differs from repeatedly calling one retained module. There is no optimizer or training loop here.

When JIT-compiling the entrypoint, make seed, rate and mode static: the wrapper converts them to Python values, and the module uses Python branches. Only x needs to be dynamic. Pass the deterministic flag explicitly to this custom module; calling model.eval() does not supply its required call argument.

Check key consumption as well as numerical output. An implementation could return the right zero-rate values while still advancing its RNG, which would change later samples on a retained module.