JAX in Action: Neural Networks from Scratch

This course guides you through building a simple Multi-Layer Perceptron (MLP) from scratch using only JAX and its NumPy API. You'll tackle the classic XOR problem, learning how to manage model parameters as PyTrees, implement the forward pass, define a loss function, compute gradients, and write a manual training loop with a basic optimizer step.