High Level Neural Network Modeling with Flax

In our previous article, we broke down why JAX requires explicit state handling. We saw that wrapping our parameters, optimizer momentum, and PRNG keys into a clean carry pattern lets us use jax.jit, jax.grad, and jax.lax.scan without running into hidden side effects.

Manually writing parameter dictionaries ({'w': ..., 'b': ...}) and manually propagating shapes works fine for a two-layer linear regression. But as soon as you build multi-layer perceptrons, convolutional networks or architectures with batch normalization and dropout, handwriting nested dictionary updates becomes tedious and brittle and this is where Flax comes in.

Flax is a high-level neural network library built natively for JAX. It gives you the clean modular ergonomics of object-oriented layer definitions while staying 100% faithful to JAX's functional and stateless philosophy under the hood.

In this post, we will explore how Flax bridges the gap between ergonomic layer design and pure functional execution, how parameter initialization works, and how to use Flax's TrainState abstraction to run a complete training workflow.

Modules as Blueprints, Not State Containers

If you are coming from PyTorch's nn.Module, you are used to a class that holds both the computational architecture and the actual weight tensors inside self as expressed in the code example below

# In PyTorch Architecture and State are tightly coupled in the object
class PyTorchMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)  # Weights live inside this attribute
        
    def forward(self, x):
        return torch.relu(self.fc1(x))

In Flax, a linen.Module does not store your parameter tensors inside the instance. Instead, the module is purely a computational blueprint. It defines how inputs should be transformed when parameters are provided.

A Flax module gives you two primary methods:

• module.init(rng_key, dummy_input) which traces the flax module blueprint using a dummy input to allocate and return an immutable dictionary of initialized parameters.

• module.apply(params, input) which takes an explicit parameter dictionary and an input array, runs the forward pass, and returns the result.

Defining Neural Networks with flax.linen

Let's consider a Multi-Layer Perceptron (MLP) with dropout and batch normalization to see how Flax structures layer definitions

import jax
import jax.numpy as jnp
from flax import linen as nn

class ClassifierMLP(nn.Module):
    hidden_dim: int
    num_classes: int
    dropout_rate: float = 0.1

    @nn.compact
    def __call__(self, x:jnp.ndarray, training: bool = False) -> jnp.ndarray:
        """
        Forward pass blue print
        """
        # Layer 1
        # Dense -> BatchNorm -> Relu -> Dropout
        x = nn.Dense(features=self.hidden_dim)(x)
        x = nn.BatchNorm(use_batch_average=not training)(x)
        x = nn.relu(x)
        x = nn.Dropout(rate=self.dropout_rate, deterministic=not training)(x)

        # Layer 2
        # Output Projection
        logits = nn.Dense(features=self.num_classes)(x)
        return logits
        

Before explaining what's going on in the code snippet above, I will draw your attention to the decorator @nn.compact. In standard Python classes, you define the child layers in the __init__ and use them in forward. However, in Flax, the @nn.compact decorator lets you define child layers inside the __call__ method right where they are used.

As for how the layers are working, the nn.Dense(features=self.hidden_dim)(x) linearly projects the input x to the hidden dimension. The nn.BatchNorm(use_running_average=not training)(x) normalizes activations accross the batch. During training when training=True, it calculates the batch statistics and updates running averages. During inference when training=False, it uses the stored running averages. The nnx.relu(x) function applies the standard ReLU activation, max(0, x). The nn.Dropout(rate=self.dropout_rate, deterministic=not training)(x) randomly drops units to prevent overfitting and finally the nn.Dense(features=self.num_classes)(x) linearly maps hidden features to the final class logits.

Initializing and Applying Modules

Because the module is just a blue print, we instantiate the class without passing any tensor data, then initialize its parameters using a PRNG key and dummy input shape.

# Instantiate the blue print
model = ClassifierMLP(hidden_dim=64, num_classes=10, dropout_rate=0.2)

# Create PRNG keys for parameter init and dropout
main_key = jax.random.key(420)
main_key, init_key, dropout_key = jax.random.split(main_key, 3)

# Create dummy batch matching input shape, batch_size=4, input_features=32
dummy_x = jnp.ones((4, 32))

# Initialize parameters and collection variables
variables = model.init(
    {'params': init_key, 'dropout': dropout_key},
    dummy_x,
    training=False
)

# Inspect the returned pytrees
print("Initialized Variable Collections:", variables.keys() )

Output

Initialized Variable Collections: dict_keys(['params', 'batch_stats'])

Notice that the model.init returned a nested dictionary containing two collections i.e. 'params', the learnable weights and biases and 'batch_stats', the running mean and variance tracked by BatchNorm.

To run an inference pass, we pass those variables back to model.apply. Please note that this is a pure and stateless pass.

# Forward inference pass 
logits = model.apply(
    variables,
    dummy_x,
    training=False
)

print("Inference Output Shape: ", logits.shape)

Output

Inference Output Shape:  (4, 10)

Structuring the Training Loop using TrainState

In our previous article, we managed our parameter dictionary and momentum state manually inside a custom carry tuple. In real applications, tracking parameters, optimizer states, batch stats and update rules manually creates boilerplate.

Flax solves this with flax.training.train_state.TrainState.

TrainState is a clean dataclass that bundles your model parameters, optimizer step and gradient update logic into a single immutable JAX PyTree.

Let's integrated Flax with Optax, a standard functional optimization library for JAX to construct a complete training step.

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.training import train_state
import optax 
from typing import Any

# Define custom TrainState to hold BatchNorm tracking statistics
class CustomTrainState(train_state.TrainState):
    batch_stats: Any

# Setup model, opitimizer and initial state
model = ClassifierMLP(hidden_dim=64, num_classes=10, dropout_rate=0.1)

init_rng = jax.random.key(42)
init_rng, init_param_key, init_drop_key = jax.random.split(init_rng, 3)

dummy_inputs = jnp.ones((8, 32))
init_variables = model.init(
    {'params': init_param_key, 'dropout': init_drop_key},
    dummy_inputs,
    training=False
)

# Setup Adam optimizer using Optax with clipping to avoid exploding gradients 
tx = optax.chain(
    optax.adam(learning_rate=1e-3),
    optax.clip_by_global_norm(1.0) 
)

# Bundle into TrainState
state = CustomTrainState.create(
    apply_fn=model.apply,
    params=init_variables['params'],
    batch_stats=init_variables['batch_stats'],
    tx=tx,
)

The JIT-Compiled Training Step

We now write a pure JIT-compiled train_step. It takes the current state, a data batch, and a PRNG key for dropout, performs the forward/backward pass, updates running statistics, and returns the updated state.

@jax.jit
def train_step(
    state: CustomTrainState,
    batch: tuple[jnp.ndarray, jnp.ndarray],
    dropout_key: jax.Array
) -> tuple[CustomTrainState, jnp.ndarray]:
    """
    Executes a forward + backward + optimizer update step.
    """
    images, labels = batch

    def loss_fn(params):
        # run forward pass passing both params and batch_stats
        # mutable=['batch_stats'] tells Flax to record updated running stats
        logits, mutated_vars = state.apply_fn(
            {'params': params, 'batch_stats': state.batch_stats},
            images,
            training=True,
            rngs={'dropout': dropout_key},
            mutable=['batch_stats']            
        )

        # Cross-entropy loss
        one_hot_labels = jax.nn.one_hot(labels, num_classes=10)
        loss = optax.softmax_cross_entropy(logits=logits, labels=one_hot_labels).mean()

        return loss, (logits, mutated_vars['batch_stats'])

    # Compute loss and gradients with respect to params
    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
    (loss_val, (logits, new_batch_stats)), grads = grad_fn(state.params)

    # Optax optimizer step that updates params and opt_state inside TrainState
    new_state = state.apply_gradients(
        grads=grads,
        batch_stats=new_batch_stats
    )

    return new_state, loss_val

Running the Training Loop

With train_step compiled, our execution loop remains clean, readable and fast

# Generate a synthetic dataset for demo purposes
dataset_size = 64
batch_size = 8
features = 32

rng_data = jax.random.key(4747)
x_data = jax.random.normal(rng_data, (dataset_size, features))
y_data = jax.random.randint(rng_data, (dataset_size,), minval=0, maxval=10)

step_key = jax.random.key(7)

# Run mini-batch training updates
num_epochs=3
for epoch in range(num_epochs):
    epochs_loss = 0.0
    num_batches = dataset_size

    for i in range(num_batches):
        step_key, subkey = jax.random.split(step_key)
        batch = (
            x_data[i * batch_size : (i + 1) * batch_size],
            y_data[i * batch_size : (i + 1) * batch_size]
        )
        # State in -> State out
        state, loss = train_step(state, batch, subkey)
        epochs_loss += loss
    print(f"Epoch {epoch + 1} | Average Loss: {epoch_loss / num_batches:.4f}")

Outputs

Epoch 1 | Average Loss: 2.7051
Epoch 2 | Average Loss: 2.5244
Epoch 3 | Average Loss: 2.4118

Notice how Flax retains the raw performance and clarity of JAX like how state remains an explicit and immutable object. The train_step is placed and compiled via @jax.jit without hidden side effects. Auxiliary states like batch_stats are tracked and updated deterministically without mutating hidden class properties.

All in all

• Flax linen.Module instances describe computation rather than storing tensor state. Parameters live separately in dictionaries returned by .init().

• Use @nn.compact to define layers and operations directly in the execution path, allowing Flax to automatically manage shape inference.

• Flax's TrainState brings together parameters, optimizer states, step counts, and custom statistics into a single immutable PyTree.

• Flax structures its APIs to ensure your training steps integrate directly with jax.jit, jax.grad, and Optax optimization chains.

Please note that the code above is better explored in a Jupyter Notebook.

See you guys in the next one.