JAX 201: performance and scaling

The JAX 101 docs covered how to express computations: arrays and jax.numpy, transformations like jax.grad() and jax.vmap(), pytrees, randomness, and state. None of that made anything fast. These pages are about performance: compiling your code, measuring it, scaling it from one chip to thousands, and tuning it.

The core arc, meant to be read in order, runs from compiling on one device to programming many:

1. Just-in-time compilation — how jax.jit() works and how to use it well: tracing and retracing, static arguments, caching, in-place updates with refs and buffer donation, and asynchronous dispatch.

2. Ahead-of-time lowering and compilation — inspecting or controlling each stage of the jit pipeline.

3. Control flow and logical operators with jit — expressing conditionals and loops so they can be compiled: constraints on Python control flow under jit, and the structured alternatives like lax.cond and lax.scan.

4. Data placement — where arrays live: meshes as the unit of placement (single-device meshes included), committed vs. uncommitted arrays, and moving data between meshes.

5. Distributed arrays and automatic parallelization — the global-view programming model that scales one program to many devices, sharding as distributed data layout, plus device-local layout.

6. Manual parallelism with shard_map — the full tutorial for sharding’s manual mode: per-device programming with explicit collectives, for complete control over how computation and communication are partitioned.

7. External callbacks — calling back to host Python from compiled code with pure_callback and io_callback, including how callbacks interact with sharded data.

Diagnostics and tuning pages are there for when you need them, in any order:

8. Benchmarking and profiling — how to measure: benchmarking pitfalls, capturing and reading profiler traces (including of distributed code), and device memory profiling.

9. Debugging runtime values — printing and inspecting values inside compiled code, and the debugging flags every JAX user should know.

10. Debugging slow JAX tracing and XLA compilation — diagnostic flags, reading the logs, and the Python patterns that defeat JAX’s caches.

11. Matmul precision — dot algorithms, the classic Precision levels, and global defaults.

12. Controlling XLA from JAX — XLA flags per function or process-wide, and attaching metadata to operations.

13. GPU memory allocation — how the allocator works, and what to do about out-of-memory failures.

14. Memory spaces and host offloading — host offloading: parking parameters, activations, and optimizer state in host memory to save device memory.

One more performance topic — computation/communication overlap — will be added here as that part of JAX stabilizes.