nnx_rngs

Chapter 3: nnx.Rngs

In the previous chapter, we explored how nnx.Variable and nnx.VariableState manage state within nnx.Module instances. We saw specific variable types like nnx.Param and nnx.BatchStat. This chapter focuses on another crucial aspect of neural network initialization and certain layer types: managing randomness using nnx.Rngs.

Motivation: Explicit Randomness in JAX

JAX's functional programming paradigm requires explicit handling of Pseudo-Random Number Generator (PRNG) keys. Unlike stateful libraries where randomness might be implicit via a global seed, JAX functions need PRNG keys passed as arguments. Manually managing and splitting these keys across complex nested modules can become tedious and error-prone.

Flax NNX introduces nnx.Rngs to streamline PRNG key management. It acts as a central container for different "streams" of randomness (e.g., one for parameter initialization, another for dropout). Modules request keys from specific streams as needed, and nnx.Rngs ensures unique keys are generated sequentially, maintaining reproducibility and fitting within JAX's functional requirements.

Central Use Case: Initializing Modules with Randomness

Most neural networks require random initialization for parameters, and some layers (like dropout) require randomness during the forward pass. nnx.Rngs is typically instantiated once at the top level and passed down the module hierarchy during construction.

import jax
import jax.numpy as jnp
from flax import nnx
from flax.nnx.nn import Linear, Dropout # Using pre-built NNX layers

class SimpleMLP(nnx.Module):
  def __init__(self, din: int, dhidden: int, dout: int, *, rngs: nnx.Rngs):
    # Pass the Rngs instance down to submodules
    self.linear1 = Linear(din, dhidden, rngs=rngs)
    # Dropout requires an RNG key for its mask generation setup
    self.dropout = Dropout(rate=0.5, rngs=rngs) 
    self.linear2 = Linear(dhidden, dout, rngs=rngs)

  def __call__(self, x: jax.Array, *, train: bool) -> jax.Array:
    x = self.linear1(x)
    x = nnx.relu(x)
    # Pass deterministic flag to control dropout behavior
    x = self.dropout(x, deterministic=not train) 
    x = self.linear2(x)
    return x

# 1. Instantiate Rngs at the top level with a base seed (0)
#    Specify different seeds for specific streams if needed.
top_level_rngs = nnx.Rngs(default=0, params=1, dropout=2)

# 2. Pass the Rngs instance during MLP initialization
mlp = SimpleMLP(din=10, dhidden=20, dout=5, rngs=top_level_rngs)

# Create dummy data
x = jnp.ones((1, 10))

# Run in training mode (dropout active)
# Dropout internally uses its 'dropout' stream key, implicitly managed by nnx.Dropout
y_train = mlp(x, train=True)

# Run in evaluation mode (dropout inactive)
y_eval = mlp(x, train=False)

print(f"MLP Output shape (train): {y_train.shape}")
print(f"MLP Output shape (eval): {y_eval.shape}")

# Inspect the state to see the RngState variables managed internally
mlp_state = nnx.state(mlp)
print("\nMLP State showing RngState for Dropout:")
# Note: Linear layers use 'params' stream *during init*, but don't store RngState after.
# Dropout stores RngState to generate masks *during calls*.
print(mlp_state['dropout'])