Learning JAX: JAX and Flax NNX Tutorials

This is a complete collection of JAX and Flax NNX tutorials. Work through them to build high-performance ML models with the JAX ecosystem.

The material comes from Google’s Learning-JAX project and covers topics from fundamentals to advanced workflows.

✨ JAX Tutorial Series

• Introduction
• How to think in JAX
• Introduction to Flax NNX
• MNIST tutorial
• Chex Exercises: Building Robust JAX & Flax NNX Applications
• Colab Notebook: Debugging JAX & Flax NNX - Exercises
• Efficient Data Loading with Grain: Exercises for JAX/Flax NNX
• Introduction to Checkpointing with Flax NNX and Orbax
• Serving a JAX model with vLLM and GPU
• Sharding and Parallelism with JAX & Flax NNX
• Optax with Flax NNX: Exercises for PyTorch Users

About JAX

JAX is a Python library for high-performance numerical computing, ideal for machine learning research. It provides:

• Automatic differentiation via jax.grad
• JIT compilation via jax.jit
• Automatic vectorization via jax.vmap
• Parallelism with straightforward distributed training

Key Components

• JAX: core numerical library
• Flax NNX: modern neural network library
• Optax: optimizer library
• Orbax: checkpointing library
• Grain: data loading toolkit
• Chex: testing and reliability tools

Resources

• JAX AI Stack
• JAX docs
• Flax docs
• Optax docs
• Orbax docs