JAX Pathways: illustrated course reader

What you will learn

Build a mental model of JAX arrays, automatic differentiation and functional transformations. Then connect Flax models, optimization, experiment tracking, recovery, evaluation and deployment. Worked equations, explained plots and Python examples connect the ideas to practical projects.

Choose a reading format

The HTML reader embeds figures and math fonts for offline reading and browser printing to PDF. The EPUB includes chapters, practical appendices and a separate solutions section. Both come from the same lessons as the online course.

For executable code and a place to keep your own results, use the downloadable project workspaces. Reference CPU runs are recorded separately from physical TPU qualification and expert review, which remain outstanding.

Chapters and online lessons

1. Setup & first steps

   Run a saved program and explain its environment, array output and device placement. Start on one CPU, practice sharding on four logical CPU devices, and learn how to get a TPU running, launch a TPU job, and run course experiments on TPU in parallel from Phase 00 onward.

2. Arrays & pure functions

   Treat every axis as part of a contract. Work through normalization, broadcasting and immutable updates, then diagnose mistakes that still return plausible numbers.

3. JAX transformations

   Begin with a scalar function, then differentiate it, batch it and compile it. Compare each transformed result with a simple reference so you can explain what changed.

4. State, randomness & control flow

   Make randomness, parameters and loop state explicit. Learn how to repeat an experiment without accidentally reusing randomness or mutating a caller’s state.

5. Math & optimization

   Connect geometry and probability to the updates in a training loop. Begin with linear algebra, gradient checks and gradient descent, then investigate conditioning, curvature, likelihoods, regularization and optimizer behavior.

6. Neural network training

   Build the same ideas into small neural networks. Understand hidden layers, parameter state, stable objectives and evaluation before moving to a small image classifier.

7. Data & checkpoint recovery

   Follow examples from the dataset into a training update, then save enough information to reproduce the next update after interruption. Track experiment configuration, data identity and metric histories with MLflow alongside complete checkpoints.

8. Transformers & language models

   Build attention and causal masking from small arrays, then connect tokenization, mixed precision, training and generation.

9. Performance diagnosis

   Measure completed work, inspect a real profiler trace and relate compiler output to memory traffic. Compare rematerialization choices while separating saved residuals, compiler memory estimates and observed latency.

10. Distributed training

    Partition arrays and training updates, compare all-reduce with reduce-scatter, and restore a complete sharded state across an epoch boundary. Verify CPU placement and arithmetic before separate cluster validation.

11. Scientific computing

    Build and differentiate a physical simulator, recover a positive rate from observations, and separate numerical accuracy, identifiability and measured execution cost.

12. Probabilistic modeling

    Specify probability models, derive exact Bayesian regression, implement and diagnose HMC, and compare variational approximations with predictive evidence.

13. Reinforcement learning

    Build a functional environment, collect randomized trajectories, train policy-gradient and clipped-policy agents, and evaluate independent seeds.

14. Pallas kernels

    Write actual Pallas grids and buffered pipelines, verify boundaries and precision on CPU, and prepare guarded target measurements before claiming an optimization win.

15. Autodiff & JAX internals

    Trace functions into jaxpr, check forward and reverse sensitivities, write exact custom derivatives and compile a small transformation with explicit supported primitives.

16. Deployment, interoperability & edge AI

    Adapt a trained model, verify framework and export contracts, compare precision policies and measure local inference before planning server or edge capacity. Package a model service in a reproducible container and verify readiness, artifact identity and request contracts.

17. Workload operations

    Launch and observe real local jobs, detect failures and stalls, restore complete training state, verify artifact quality and rehearse rollback before planning capacity. Connect MLOps data/release gates, LLMOps evaluation/traces and ModelOps ownership/approval to the workload runbook.

18. Self-supervised pretraining: masked and contrastive learning

    Train small masked-token, hidden-patch and paired-view models. Audit information leakage, reduction counts and pair identity before interpreting a falling objective.

19. Post-training: SFT, LoRA, reward models and RLHF

    Connect supervised demonstrations, frozen-base LoRA, reward modeling, sampled PPO and direct preference optimization. Keep the reference, behavior policy, reward model and dataset identities separate.

20. TPU Setup: Provisioning & Runtime Verification

    Step 1 of the 3-part Cloud TPU mini-course. Choose the right TPU slice, provision a Cloud TPU VM (or Colab TPU), bootstrap ~/jax-tpu-lab/.venv with jax[tpu], verify a synchronized pmap runtime receipt, and practice clean teardown.

21. TPU Workflows: Launching Jobs & Checkpointed Experiments

    Step 2 of the 3-part Cloud TPU mini-course. Stage scripts with gcloud compute tpus tpu-vm scp, separate XLA compilation warmup from steady-state step timing, recover interrupted runs from atomic checkpoints, and launch single-host or --worker=all Pod slice experiments.

22. TPU Systems: Generations, Memory, Precision & XProf

    Step 3 of the 3-part Cloud TPU mini-course. Compare TPU v5e, v5p, and v6e specifications and ridge points, budget training and KV-cache HBM before provisioning, measure BF16 and INT8 numerical error, and capture warmed XProf traces.

Learn with a project or a route

Choose a learning path, explore career routes and evidence, or connect the chapters through staged JAX projects.

Found a mistake? Open a GitHub issue. Read the edition history to see what changed.