JAX in Action: Building an Image Classifier

Apply your JAX, Flax, and Optax skills to a practical project: building an image classification pipeline. This course covers setting up a multi-file project, loading and preprocessing image data (e.g., MNIST), defining a Convolutional Neural Network (CNN) with Flax, and implementing robust training and evaluation loops.