XOR Revisited: Building a Full Training Loop with Flax and Optax

Welcome back to Beyond Pure JAX: Flax & Optax for Elegant ML! In the previous lessons, you've mastered the essential building blocks: Flax modules for elegant neural network construction, powerful MLP architectures with dense layers, and sophisticated Optax optimizers with JIT-compiled training steps. Today, for the final lesson of the course we're bringing everything together to solve the classic XOR problem using our complete Flax and Optax toolkit.

By the end of this lesson, you'll have assembled a complete, production-ready training pipeline that seamlessly integrates model initialization, optimization state management, training loops, and model evaluation. This represents the culmination of everything we've learned and serves as a template for tackling real-world machine learning problems. We'll also compare our elegant Flax and Optax solution with the manual JAX implementation we built back in Course 3. Get ready!

Creating a robust training pipeline requires orchestrating several components that we've explored individually. We need to coordinate our Flax MLP model, the Optax Adam optimizer, and the JIT-compiled training step we mastered previously. The beauty of this integration lies in how these components work together seamlessly, despite being developed independently.

Our training pipeline follows a clear structure: initialize the model and create its parameters, set up the optimizer and its internal state, then iterate through training epochs while calling our compiled training step function. Each component maintains its own state that gets passed through the training loop, ensuring we can track both parameter evolution and optimizer momentum terms throughout the learning process.

The key insight is that modern deep learning frameworks like Flax and Optax are designed for composability. We can mix and match different model architectures, loss functions, and optimizers without rewriting our core training logic. This modularity makes our code both more maintainable and more experiment-friendly, allowing us to rapidly prototype different configurations.

Let's begin by setting up our core components. We'll initialize our MLP model with the same architecture we used previously, then create our optimizer with carefully chosen hyperparameters:

# Define training parameters
num_epochs = 2000
learning_rate = 0.05
prng_key = jax.random.key(123)

# Initialize the Flax MLP model
xor_mlp_model = MLP(hidden_features=3, output_features=1)

# Initialize model parameters
initial_variables = xor_mlp_model.init(prng_key, xor_X)
current_params = initial_variables['params']

print("Initial MLP Parameters (shapes):")
print(jax.tree_util.tree_map(lambda x: x.shape, current_params))

This code outputs:

Initial MLP Parameters (shapes):
{'HiddenLayer': {'bias': (3,), 'kernel': (2, 3)}, 'OutputLayer': {'bias': (1,), 'kernel': (3, 1)}}

The model.init() call is where Flax performs its lazy initialization magic. By passing our XOR input data, Flax can infer all the necessary shapes and create appropriately sized weight matrices and bias vectors. The returned initial_variables contains our parameter tree under the 'params' key, which we'll use throughout training. The jax.tree_util.tree_map function elegantly displays the shape of each parameter in our nested structure without printing the actual values.