# High Level Neural Network Modeling with Flax

In our [previous article](https://marshud.dev/functional-state-management-and-low-level-control-flow-in-jax), 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

```python
# 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.

![](https://cdn.hashnode.com/uploads/covers/68874bca8a40b70ffef830eb/95848d72-cde1-4ff6-8711-53811cc01f75.png align="center")

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

```python
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.

```python
# 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

```plaintext
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.

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

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

Output

```plaintext
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.

```python
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`.

```python
@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

```python
# 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

```plaintext
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.