Skip to content

← Flax NNX step 9 of 100

Medium Primitives

NNX Train/Eval Flag

Jump to the task

Why this matters

A layer may behave differently during training and evaluation. This exercise defines a custom layer that adds Gaussian noise only when its train argument is true. In evaluation mode it returns the input unchanged and does not draw a random key.

This is additive noise, not dropout. Dropout masks activations and rescales retained values during training; it does not add Gaussian noise.

A complete example

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

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

    def __call__(self, x, train: bool):
        if train:
            key = self.rngs.dropout()
            noise = jax.random.normal(key, x.shape, dtype=x.dtype)
            return x + 0.1 * noise
        return x

model = NoiseLayer(nnx.Rngs(0))
x = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
eval_output = model(x, train=False)
first_training_output = model(x, train=True)
second_training_output = model(x, train=True)
assert jnp.array_equal(eval_output, x)
assert int(model.rngs.default.count[...]) == 2

The RNG container keeps a counter. Each training call requests a new key; evaluation leaves the counter alone. Fresh keys produce new draws, but do not imply a mathematical guarantee that every sampled value differs.

With nnx.Rngs(0), requesting dropout() falls back to the default stream. The name does not turn the operation into dropout or create a separate stream. To configure an explicit stream, use a named seed such as nnx.Rngs(dropout=0). Use the default-stream constructor required below to match the fixtures.

Custom flags and built-in evaluation settings

train is an argument chosen by this custom layer, not a universal NNX calling convention. Built-in Dropout uses deterministic; BatchNorm uses use_running_average. NNX also provides model.train() and model.eval() to change corresponding attributes on supported nested modules.

Those methods do not supply this custom train argument automatically. Pass it explicitly. A layer can also be stochastic without having any train/eval flag at all; its behavior depends on its implementation.

The Python if train requires a concrete Boolean during tracing. For a JIT-compiled call using this implementation, make train static. Supporting a dynamic array-valued condition needs suitable JAX/NNX control flow instead of converting a tracer to a Python bool.

Problem

Implement train_eval_forward(seed, x, train):

  1. Convert the supplied mode scalar with train_b = bool(train >= 0.5).
  2. Define NoiseLayer as above, storing the RNG container on the module.
  3. Construct a fresh nnx.Rngs(int(seed)) and NoiseLayer.
  4. Return model(x, train=train_b).

The wrapper constructs a fresh model for each invocation. Repeating the entire entrypoint with identical inputs therefore repeats the first draw. This differs from repeatedly calling one persistent NoiseLayer, as in the example, where the RNG counter advances.

For training, return x + 0.1 * noise, where noise has independent standard normal entries and the same shape and float32 dtype as x. Do not request a key in the constructor or evaluation branch. If JIT-compiling the whole entrypoint, treat both seed and train as static arguments.

Inputs: seed is a nonnegative integer-valued scalar supplied as a float; cast it to int. x is a nonempty finite float32 vector. train is either 0.0 for evaluation or 1.0 for training.

Output: an array with the same shape as x, unchanged in evaluation mode and with the specified additive noise in training mode.