Core learning path (first pass)

119 must-read extras in total; 243 optional sections can be skipped on the first pass.

How to use it: work through the main text of each notebook in order (every unmarked explanation and code cell is part of the main path), plus the 🟢 Must-read extras listed below. Sections marked ⚪ Optional can wait until you need them. Try at least the first two exercises in each notebook: write your code in ✍️ Your answer, check it with 🧪 Self-check, and find reference solutions in the appendix at the end of the chapter.

Stage 1: Basics

basics/01_introduction　(3 must-read · 4 optional to skip)

• 🔍 Extra (1.2): Three more behavior differences that trip people up
• 🔍 Extra (1.3)
• 🔍 Extra (1.5-a): The number one random-number mistake: reusing a key

basics/02_jit　(5 must-read · 9 optional to skip)

• 🔍 Extra: Ways to write jit, and a very common trap
• 🔍 Extra: Prefer jnp.where for element-wise conditions, but watch out for NaN in gradients
• 🔍 Extra: Hard requirements and details for static arguments
• 🔍 Extra: Use jax.debug.print for debug printing and jax.disable_jit() for debugging logic
• 🔍 Extra: Old and new values of a global coexisting, and passing large arrays as arguments too

basics/03_grad　(4 must-read · 7 optional to skip)

• 🔍 Extra: look at the jaxpr yourself, and how jvp/vjp relate to the two modes
• 🔍 Extra: has_aux=True for functions that return more than the loss
• 🔍 Extra: the return structure when combined with has_aux
• 🔍 Extra (special topic): NaN gradients and the "double where" trick

basics/04_vmap　(7 must-read · 13 optional to skip)

• 🔍 Extra: where manual vectorization most often goes wrong
• 🔍 Extra: Python print vs. jax.debug.print
• 🔍 Extra: no Python if on batched data inside vmap
• 🔍 Special topic: per-example gradients
• 🔍 Extra: vmap and random numbers, mapping over keys
• 🔍 Extra (measured in this environment): pass large arrays as arguments; don't let a jit function capture them in a closure
• 🔍 Extra: under batching, lax.cond becomes select, and both branches run

basics/05_pmap　(2 must-read · 3 optional to skip)

• 🔍 Extra: the actual status of pmap in JAX 0.12
• ⚠️ Special topic: gradients become N times larger in shard_map (the top pmap → shard_map migration pitfall)

basics/06_state　(4 must-read · 6 optional to skip)

• 🔍 Extra: the "runs only once during tracing" trap when modifying/reading globals inside jit
• 🔍 Extra: lax.scan's carry is the state; record a trajectory with per-step outputs
• 🔍 Extra: a random key is itself a piece of state that must be passed explicitly
• 🔍 Extra: use out_axes to avoid needlessly copying shared state

basics/07_practical_topics　(16 must-read · 41 optional to skip)

• Topic 1 Recompilation: triggers, diagnosis, and fixes
• 1.2 How it works: what's in the jit cache key
• 1.5 ✅ Fixes
• 💡 Practical rule
• 2.2 Double where: how it works and a general pattern
• 2.3 Stable logsumexp / softmax / log_softmax
• 2.4 Locating NaNs with jax_debug_nans (measured behavior in 0.12)
• Topic 3 Choosing between jacfwd / jacrev / jvp / vjp
• 3.4 ❌ Common mistake: building $J$ when you only need $Jv$
• 4.3 ❌ Common mistakes
• 5.4 ❌ Common mistakes, measured
• 6.5 ❌ Common mistakes, measured
• 6.6 ✅ A standard template for the training step
• 7.4 ❌ Common mistakes, measured
• 7.5 TrainState: register_dataclass vs. NamedTuple
• ✔ Practical rules checklist

Stage 2: Debugging and random numbers (Advanced 08 and 10, right after Basics)

advanced/08_debugging　(14 must-read · 36 optional to skip)

• 0 Symptom lookup table
• 1 Seeing values: runtime numbers inside compiled code
• 1.4 jax.disable_jit(): a debugging aid and a bisection tool
• 2 Seeing the program: what JAX actually traced and compiled
• 3 Numerical problems: NaN, inf, out-of-bounds, division by zero and wrong gradients
• 3.2 jax.experimental.checkify
• 4.1 TracerBoolConversionError
• 4.5 Tracers in boolean mask indexing, shapes and slices
• 4.6 UnexpectedTracerError: tracer leaks
• 5.2 jax_explain_cache_misses: three explanations, and one case where it "goes quiet"
• 5.5 Slow code: what to check, in order
• 7.3 jax.live_arrays(): finding arrays that weren't freed
• 9.4 Hands-on recap
• Chapter summary

advanced/10_random_sharding　(12 must-read · 20 optional to skip)

• 1 Design: why there's no global state
• 4.5 "Random numbers aren't guaranteed to match across versions", and the 10.81 from Chapter 1
• 5 Key management patterns
• 5.2 split(key, n) makes several at once
• 5.3 fold_in(key, step): deriving keys for training steps / data indices
• 5.4 Putting the key in the carry of lax.scan
• 5.5 Mapping over a key array with vmap
• 6 Common pitfalls
• 6.2 Capturing a key in a closure inside jit: the random numbers become constants
• 6.4 Python random / NumPy random numbers run only once in jit, during tracing
• 8 Practical rules checklist
• Chapter summary

Stage 3: Advanced chapters 1–6

advanced/01_transformations_control_flow　(6 must-read · 10 optional to skip)

• 🔍 Extra: reading carry-type mismatch errors
• 🔍 Extra: lax.switch and lax.select_n
• 🔍 Extra: cond turns into select under vmap
• 1.6 Advanced masking techniques
• Course code: masking attention scores
• 🔍 Extra: four problems with capturing large arrays in closures

advanced/02_performance　(2 must-read · 6 optional to skip)

• 🔍 Extra
• Why v2/v3 are faster

advanced/03_distributed　(6 must-read · 11 optional to skip)

• 🔍 Extra: memory and communication budgets for the three kinds of parallelism
• Device data residency
• 🔍 Extra: three "views" of the same SPMD program
• Other collectives
• 🔍 Extra: combining a 2D mesh with PartitionSpec
• 🔍 Extra: AxisType — Explicit / Auto / Manual

advanced/04_advanced_autodiff　(6 must-read · 9 optional to skip)

• 🔍 Extra: reverse mode = linearize + transpose
• 🔍 Extra: signature details, residuals, and nondiff_argnums
• 🔍 Extra: verifying custom rules with jax.test_util.check_grads
• 🔍 Extra: JVP rules must be linear in the tangent
• 🔍 Extra: differentiability of each control-flow primitive at a glance (tested)
• 🔍 Extra: tested gradient behavior of common "non-differentiable" ops

advanced/05_interop_custom_ops　(1 must-read · 7 optional to skip)

• 🔍 Extra

advanced/06_large_scale_training　(5 must-read · 8 optional to skip)

• 🔍 Extra: how many bytes does each parameter really need? And how big are activations?
• 🔍 Extra: a few "new rules" for NNX in Flax 0.12
• 🔍 Extra: when is accumulation not equal to a large batch?
• 🔍 Extra: stacking layers with lax.scan + checkpointing is the standard structure for large models
• Things to keep in mind

Stage 4: Pick what you need (Advanced 07, 09, 11)

advanced/09_export_deploy　(6 must-read · 12 optional to skip)

• §2 jax.export: exporting the computation graph as a portable program
• §2.2 Serialize to a file and run it in a brand-new subprocess
• §2.3 Shape polymorphism: export once, handle any batch size
• §2.6 Version compatibility: three layers of version numbers
• §3.2 CheckpointManager: multi-step saving and retention policies
• §3.4 Working with Flax NNX: save, restore, keep training

advanced/07_pallas　(9 must-read · 17 optional to skip)

• §1 Why Pallas
• §2 Programming model
• 2.2 grid, BlockSpec and pl.program_id
• 3.2 Out-of-bounds blocks: you read "garbage", which is NaN in interpret mode
• §4 Advanced kernels
• 4.3 Handling boundaries with masks
• 4.4 Reductions across the grid
• 5.3 grad: not differentiable by default; use custom_vjp + a Pallas backward kernel
• §6 Notes on GPU/TPU

advanced/11_scientific_computing　(11 must-read · 24 optional to skip)

• 1 Why use JAX for scientific computing
• 2 The finite-difference Poisson equation
• 2.3 Solving with jax.scipy.sparse.linalg.cg
• 4 Implicit differentiation of solvers
• 4.3 The two ways to differentiate: implementation and cross-check
• 4.4 symmetric and transpose_solve: the pitfall with nonsymmetric operators
• 4.5 jax.scipy.sparse.linalg.cg is itself implicit differentiation
• 5 Inverse problem: recovering the diffusion coefficient $k(x)$
• 6 Time-dependent problems: the heat equation
• 6.2 Sensitivities: gradients with respect to physical parameters and initial conditions
• 9.2 The checklist