Train vs Eval Branches via train Flag

Dropout randomly removes some values during training and leaves them unchanged during evaluation. In Flax Linen, your module can accept a train argument and pass the corresponding setting to its dropout layer.

An explicit training flag

class TrainEvalModel(nn.Module):
    @nn.compact
    def __call__(self, x, train: bool):
        x = nn.Dense(features=x.shape[-1])(x)
        return nn.Dropout(rate=0.5, deterministic=not train)(x)

train is an argument you define, not a special Linen field. Here it controls only dropout. train=True makes deterministic=False, enabling random masks. train=False makes dropout pass through the Dense output unchanged. Neither setting performs an optimizer update.

The Dense layer first computes z = x @ kernel + bias. With training enabled, each element of z is independently kept with probability 0.5. Kept values are doubled and dropped values become zero. This preserves each element's expected value over random masks. A particular mask need not drop exactly half the elements.

Initialization does not register a dropout collection

This model's variables contain only params. Dropout has no learned weights or persistent mask. The named dropout RNG stream supplies random keys; it is not a variable collection that must be created during init.

Initialization with train=False works with just the parameter key. You can then enable dropout during apply by supplying a dropout key. This exercise passes both named keys during initialization for an explicit, fixed recipe:

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

In evaluation, this model needs no rngs argument; providing one is harmless. During training, pass a fresh dropout key for each step when you want fresh randomness. Repeating the same key and call structure repeats the mask. Reusing it here makes the exercise reproducible.

Compiling the forward pass

Python's not train needs a concrete boolean. When using jax.jit on model.apply, mark train static with static_argnames=("train",) or close over a fixed training setting. Training and evaluation can then use separate compiled paths. An ordinary traced boolean cannot be used with Python not.

Other layers have their own controls. For example, BatchNorm can use and update running statistics. Do not assume that toggling this dropout flag handles every kind of training state in a larger model.