Advanced · Overview

Chapters 1–6 follow the apxml.com course Advanced JAX: Performance, Optimization, and Scale. Chapters 7–11 are extra chapters on topics that matter a lot when you actually use JAX, written in the same style as the earlier chapters. Every notebook has been run end to end on JAX 0.12.0.1 / Python 3.14 / CPU.

What to learn first: the Basics part, 00 Overview, especially 02 JIT, 03 grad, 05 Parallelism and 07 Practical topics. The Advanced chapters refer to these directly and don't repeat them.

Reading conventions (same as in Basics)

The main code uses the style JAX 0.12 recommends. Outdated styles go in # ❌ Old style cells, followed right away by # ✅ Recommended.
Things to watch out for are marked > ⚠️ Note, and important conclusions come with an assert.
The explanations reorganize the course's key points, and each section starts with a link to the original course page. Quiz questions are based on the chapter content.
Exercises: write your code in the "✍️ Your answer" cell below each exercise and run the "🧪 Self-check" to check it. Reference solutions are in the appendix at the end of each chapter.

Chapter index¶

Ch.	Notebook	Highlights	Version notes
1	01 Advanced transformations and control flow	scan argument details and `associative_scan`; masking and `segment_*`; closures and constvars; hands-on: variable-length masked LSTM, RK4 ODE integrator	22 items
2	02 Optimizing JAX code performance	Real profiler traces on CPU; the full jaxpr → StableHLO → HLO pipeline; persistent compilation cache (cross-process hits measured); quantifying fusion with HLO (memory traffic 3N vs 7N); tracking down the root cause of a CPU bug	19 items
3	03 Distributed computing in JAX	A modern-style version of every pmap example; exact semantics of collective operations; Megatron MLP; FSDP per-device memory 26144 B → 5956 B; a real 2-process multi-host demo on one machine	26 items
4	04 Advanced automatic differentiation	`linearize` + `linear_transpose`; higher-order Taylor expansion with `jet`; Jacobian coloring; implicit function theorem and custom_vjp; DEQ implicit layers, ODE adjoint	18 items
5	05 Interoperability and custom operations	Measured comparison table of the three callbacks; DLPack zero-copy; the full set of rules for a custom primitive (jit/grad/vmap); compiling a C++ kernel with clang++ and plugging it in through `jax.ffi`	29 items
6	06 Large-scale model training	Flax NNX + Optax + Orbax; three non-equivalent cases of gradient accumulation; quantitative comparison of checkpoint policies (Transformer temporary memory down 84%); bf16/fp16 and loss scaling; tensor parallelism, FSDP, Adafactor	19 items

Extra chapters

Ch.	Notebook	Contents	Read first
7	07 Pallas custom kernels	The Pallas programming model (Ref, grid, BlockSpec); checking correctness on CPU in interpret mode; blocked matmul, softmax, LayerNorm, masking; differentiating with custom_vjp; cross-lowering to TPU to check block shape rules	Chapter 2 (fusion boundaries), Chapter 5 (custom primitives / FFI)
8	08 Debugging toolchain	A troubleshooting manual organized as "symptom → tool": debug.print / callback, jaxpr and HLO, debug_nans and checkify, the various Tracer errors, compilation diagnostics, memory investigation; hands-on: fixing a training script with 3 hidden bugs	Basics 02, 07; Chapter 2
9	09 Export and deployment	Parameters with Orbax, programs with `jax.export`; shape polymorphism, multiple platforms, `vjp_order`, version compatibility; running inference in a fresh subprocess that loads only the exported file and the parameters	Chapter 2 (compilation cache), Chapter 6 (Orbax)
10	10 Random numbers and sharding	Typed keys vs raw keys; 6 PRNG implementations; how `jax_threefry_partitionable` affects sharded generation and reproducibility across versions; key management patterns and common pitfalls	Basics 01, 07; Chapter 3
11	11 Scientific computing and differentiable simulation	float64; matrix-free Poisson solver (second-order convergence verified); BCOO sparse matrices; implicit differentiation through a solver (about 290x less memory); a coefficient inverse problem; sensitivity analysis for the heat equation; Winslow mesh optimization, inverted elements down from 23 to 0	Chapter 1 (scan / ODE), Chapter 4 (implicit function theorem)

Suggested order for the extra chapters: 08 Debugging toolchain (used the most; you can read it right after Basics) → 10 Random numbers (it also explains why the example outputs in Basics Chapter 1 can't be reproduced) → 09 Export and deployment → 07 Pallas (when you need to hand-write high-performance kernels) → 11 Scientific computing (most valuable if you work on numerical simulation or meshes).

Chapters 3, 6 and 10 set XLA_FLAGS at the top to simulate 8 devices, so restart the kernel and run from the first cell. Chapter 5 needs a C++ compiler on your machine (the clang++ that ships with macOS is fine). Serialization in Chapter 9 needs flatbuffers (already in requirements.txt). Chapter 11 uses float64 throughout.

Environment check¶

import sys, shutil, importlib
import jax
print("Python :", sys.version.split()[0])
print("JAX    :", jax.__version__)
for name in ["optax", "flax", "orbax.checkpoint", "flatbuffers"]:
    try:
        m = importlib.import_module(name)
        print(f"{name:7s}:", getattr(m, "__version__", "installed"))
    except ImportError:
        print(f"{name:7s}: not installed (needed for Chapters 6 and 9)")
print("clang++:", shutil.which("clang++") or shutil.which("c++") or "not found (needed for the C++ FFI in Chapter 5)")
print("Devices:", jax.devices())

Python : 3.14.7
JAX    : 0.12.0.1
optax  : 0.2.9
flax   : 0.12.12

orbax.checkpoint: 0.12.10
flatbuffers: 25.12.19
clang++: /usr/bin/clang++
Devices: [CpuDevice(id=0)]

Key points of the Advanced chapters¶

1. Removed or outdated APIs (as of 0.12)

Old style	What happens in 0.12	Current style	Ch.
`jax.experimental.host_callback`	Module removed	`jax.debug.callback` / `jax.experimental.io_callback` / `jax.pure_callback`	5
`jax.dlpack.to_dlpack`, bare DLPack capsules	Removed / TypeError	Just `jax.dlpack.from_dlpack(obj)`, `np.from_dlpack(x)`	5
`jax.core.Primitive`	Doesn't exist	`jax.extend.core.Primitive`	5
`xla_client.register_translation`, `CustomCall`	Not available	`mlir.register_lowering`; for C++, use the XLA FFI (`jax.ffi`)	5
`pure_callback(..., vectorized=True)`	The argument is passed to the callback as a kwarg	`vmap_method="sequential"` / `"expand_dims"`, etc.	5
`device_put_replicated` / `device_put_sharded`	AttributeError	`jax.device_put(x, NamedSharding(mesh, P(...)))`	3, 6
Nested pmap, `jax.experimental.maps`	Error / removed	2D `jax.make_mesh` + `jax.shard_map`	3
Haiku, Flax Linen	Haiku is in maintenance mode; Linen is the old API	Flax NNX (`.value` is deprecated, use `var[...]`)	6

2. Easy things to get wrong

while_loop supports only forward-mode differentiation; reverse mode (grad) raises an error (Chapters 1, 4).
You can't take grad of a pure_callback directly; you have to supply the derivative with custom_jvp/vjp (Chapter 5).
If a constant captured by a closure changes, jit does not recompile; it keeps using the old value (Chapter 1).
Wrapping the whole model in one checkpoint doesn't lower peak memory at all; you need per-layer checkpoints for it to help (Chapter 6).
Under mixed precision, the gradient with respect to fp32 parameters is fp32, not bf16 (Chapter 6).
Optax itself doesn't handle sharding; where the optimizer state lives follows the parameters (Chapter 6).
pmean and psum/N compile to the same thing; neither one is faster (Chapter 3).

The "Version notes" section at the end of each chapter has the full list, 133 items in total.

Where to go next¶

Training infrastructure: Grain (data loading), and MaxText as a reference implementation.
More on custom kernels: run the Pallas kernels from Chapter 7 on a real GPU/TPU (see the GPU verification guide).
Scaling theory: the official How to Scale Your Model.