Keyword index

Concepts

Each concept lists its most relevant sections (heading matches first), up to 8.

Tracing and jaxpr

• Advanced 02_performance › 2.2 Understanding the JAX computation graph (jaxpr)
• Advanced 02_performance › 2.1 Profiling JAX code on CPU, GPU and TPU
• Basics 02_jit › 2.3 How JIT works: tracing and compilation
• Basics 02_jit › 2.4 Python control flow and jit
• Advanced 08_debugging › 4.6 UnexpectedTracerError: tracer leaks
• Advanced 08_debugging › 6.1 JAX_TRACEBACK_FILTERING: how many JAX-internal frames show up in a traceback
• Advanced 08_debugging › 2.1 jaxpr: spotting values that were treated as constants
• Basics 02_jit › 2.5 Static vs. traced values

Recompilation and the compilation cache

• Advanced 02_performance › 2.5 Avoiding recompilation
• Basics 02_jit › 2.3 How JIT works: tracing and compilation
• Basics 02_jit › 2.5 Static vs. traced values
• Basics 07_practical_topics › 1.1 Motivation
• Basics 07_practical_topics › ★.4 Training: compile budget = 1
• Advanced 02_performance › 2.2 Understanding the JAX computation graph (jaxpr)
• Advanced 08_debugging › 9.2 Symptom 2: recompiling every step, locating it
• Basics 02_jit › Exercise 3: recompilation detective

Static arguments and weak types

• Basics 02_jit › 2.5 Static vs. traced values
• Basics 07_practical_topics › 1.5 ✅ Fixes
• Advanced 02_performance › 2.5 Avoiding recompilation
• Basics 02_jit › 2.7 Hands-on: applying jit
• Basics 02_jit › 2.3 How JIT works: tracing and compilation
• Basics 06_state › 6.5 Example: a stateful counter
• Basics 06_state › 6.6 Example: simple optimizer state
• Basics 06_state › 6.8 Practice: implementing stateful functions

Control flow (cond / scan / while)

• Basics 02_jit › 2.4 Python control flow and jit
• Advanced 01_transformations_control_flow › 1.4 Looping with lax.while_loop
• Basics 03_grad › 3.7 Differentiation and control flow
• Basics 04_vmap › 4.7 Performance considerations for vmap
• Advanced 01_transformations_control_flow › 1.3 Conditional execution with lax.cond
• Advanced 04_advanced_autodiff › 4.8 Differentiating through control-flow primitives
• Advanced 02_performance › 2.2 Understanding the JAX computation graph (jaxpr)
• Basics 06_state › 🏋️ Extra exercises (with reference solutions)

Masking and variable-length sequences

• Advanced 01_transformations_control_flow › 1.6 Advanced masking techniques
• Basics 07_practical_topics › 1.5 ✅ Fixes
• Basics 02_jit › 2.6 Common jit pitfalls
• Advanced 07_pallas › 4.3 Handling boundaries with masks
• Advanced 01_transformations_control_flow › 1.8 Hands-on: implementing complex recurrent logic
• Basics 07_practical_topics › Exercise 6: a training step with frozen layers + a donated state
• Advanced 08_debugging › 4.5 Tracers in boolean mask indexing, shapes and slices
• Basics 07_practical_topics › 2.3 Stable logsumexp / softmax / log_softmax

NaN and numerical stability

• Basics 03_grad › 3.8 Limitations and caveats
• Basics 07_practical_topics › 2.4 Locating NaNs with jax_debug_nans (measured behavior in 0.12)
• Advanced 08_debugging › 3.1 jax_debug_nans / jax_debug_infs
• Basics 07_practical_topics › 2.1 Motivation
• Advanced 11_scientific_computing › 7.4 A numerically stable barrier: double where
• Basics 07_practical_topics › 2.3 Stable logsumexp / softmax / log_softmax
• Advanced 08_debugging › 9.1 Symptom 1: NaN, locating it
• Advanced 08_debugging › Exercise 4: automatically locating the first NaN

Autodiff modes (jvp / vjp)

• Basics 03_grad › 3.3 How autodiff works: reverse mode
• Advanced 04_advanced_autodiff › 4.3 Vector-Jacobian products (VJPs) and jax.vjp
• Basics 07_practical_topics › 3.3 Timing experiment: sweeping input/output dimensions
• Advanced 04_advanced_autodiff › 4.1 Review of forward- and reverse-mode autodiff
• Advanced 07_pallas › 5.3 grad: not differentiable by default; use custom_vjp + a Pallas backward kernel
• Advanced 09_export_deploy › §2.5 Exporting differentiable functions: vjp_order
• Basics 03_grad › 3.8 Limitations and caveats
• Advanced 04_advanced_autodiff › 4.2 Jacobian-vector products (JVPs) and jax.jvp

Jacobian / Hessian

• Advanced 04_advanced_autodiff › 4.5 Computing full Jacobians and Hessians
• Basics 03_grad › 3.3 How autodiff works: reverse mode
• Basics 07_practical_topics › 3.3 Timing experiment: sweeping input/output dimensions
• Basics 03_grad › 3.5 Higher-order derivatives (grad of grad)
• Advanced 11_scientific_computing › Exercise 4: verifying the corner Jacobians with jacfwd
• Basics 07_practical_topics › 3.5 Three ways to compute Hessian-vector products (HVP)
• Basics 03_grad › Extra exercises
• Advanced 04_advanced_autodiff › 4.2 Jacobian-vector products (JVPs) and jax.jvp

Custom derivatives (custom_jvp / custom_vjp)

• Basics 03_grad › 3.8 Limitations and caveats
• Advanced 04_advanced_autodiff › 4.9 Handling non-differentiable functions
• Advanced 07_pallas › 5.3 grad: not differentiable by default; use custom_vjp + a Pallas backward kernel
• Advanced 04_advanced_autodiff › 4.10 Practice: implementing custom gradients
• Advanced 04_advanced_autodiff › Extra exercise D: a vector norm that is safe at 0 (custom_jvp)
• Basics 07_practical_topics › 2.5 Fixing gradients with custom_jvp
• Advanced 08_debugging › 1.1 jax.debug.print
• Advanced 04_advanced_autodiff › 4.8 Differentiating through control-flow primitives

Per-example gradients

• Basics 04_vmap › 4.6 Combining vmap with jit and grad
• Basics 04_vmap › 4.8 Hands-on: vectorizing functions
• Basics 07_practical_topics › 4.2 Cross-checking vmap(grad) against a loop
• Basics 07_practical_topics › 4.5 DP-SGD: a complete implementation with clipping + noise
• Advanced 01_transformations_control_flow › 1.1 Core transformations review: jit, grad, vmap
• Basics 07_practical_topics › ★.6 Per-example gradients: finding the top 5 "hard examples"
• Basics 07_practical_topics › Exercise 4: finding mislabeled examples with per-example gradients
• Basics 04_vmap › 4.3 Mapping over specific arguments (in_axes, out_axes)

Vectorization with vmap

• Basics 04_vmap › 4.8 Hands-on: vectorizing functions
• Basics 04_vmap › 4.7 Performance considerations for vmap
• Basics 04_vmap › Sections
• Basics 04_vmap › 4.6 Combining vmap with jit and grad
• Advanced 01_transformations_control_flow › 1.1 Core transformations review: jit, grad, vmap
• Basics 04_vmap › 4.2 Introducing jax.vmap
• Basics 04_vmap › 4.5 Nested vmap
• Basics 04_vmap › 4.1 The concept of vectorization

Multiple devices and sharding

• Advanced 03_distributed › 3.7 Nested pmap and advanced partitioning
• Advanced 06_large_scale_training › 6.4 Combining pmap with training frameworks
• Basics 05_pmap › 5.4 Device meshes and axis names
• Basics 05_pmap › ⚠️ Special topic: gradients become N times larger in shard_map (the top pmap → shard_map migration pitfall)
• Advanced 06_large_scale_training › 6.8 Model parallelism strategies
• Advanced 09_export_deploy › §2.7 Exporting functions with sharding, shard_map or pmap
• Basics 05_pmap › 5.2 Introducing jax.pmap
• Advanced 08_debugging › 8 Multiple devices: where arrays live and what each device computed

pmap and collectives

• Advanced 03_distributed › 3.5 Collective communication primitives (psum, pmean, etc.)
• Basics 05_pmap › 5.5 Collective operations (lax.psum, lax.pmean, etc.)
• Basics 05_pmap › 5.2 Introducing jax.pmap
• Basics 05_pmap › 5.7 Debugging pmapped functions
• Advanced 06_large_scale_training › 6.4 Combining pmap with training frameworks
• Basics 05_pmap › 5.6 Combining pmap with other transformations
• Advanced 03_distributed › 3.7 Nested pmap and advanced partitioning
• Advanced 03_distributed › 3.6 Handling axis names in pmap

Model parallelism (FSDP / tensor parallel)

• Advanced 06_large_scale_training › 6.8 Model parallelism strategies
• Advanced 03_distributed › 3.1 Introduction to parallelism
• Advanced 03_distributed › 3.7 Nested pmap and advanced partitioning
• Advanced 03_distributed › Exercises
• Advanced 03_distributed › 3.5 Collective communication primitives (psum, pmean, etc.)
• Advanced 06_large_scale_training › 6.1 Challenges in training large models
• Basics 05_pmap › 5.8 Hands-on: parallel computation
• Advanced 06_large_scale_training › 6.9 Optimization algorithms at scale

Multi-host

• Advanced 10_random_sharding › 5.6 Multiple processes: deriving a key for each process
• Advanced 03_distributed › 3.2 Device management in JAX
• Advanced 06_large_scale_training › 🧭 Where to go next
• Basics 05_pmap › 5.8 Hands-on: parallel computation
• Advanced 09_export_deploy › §2.7 Exporting functions with sharding, shard_map or pmap
• Advanced 10_random_sharding › Exercises

PyTrees

• Basics 06_state › 6.4 Using PyTrees for structured state
• Basics 07_practical_topics › 7.2 Old → new migration table (measured in 0.12)
• Basics 06_state › 6.6 Example: simple optimizer state
• Basics 06_state › 6.7 Combining state management with transformations
• Advanced 01_transformations_control_flow › 1.2 Mastering lax.scan for sequential operations
• Basics 04_vmap › 4.3 Mapping over specific arguments (in_axes, out_axes)
• Advanced 06_large_scale_training › 6.3 Managing model parameters and state
• Basics 05_pmap › 5.3 Mapping data to devices (in_axes, out_axes)

Random numbers and keys

• Basics 01_introduction › 1.5 Working with JAX arrays
• Advanced 10_random_sharding › 3 PRNG algorithms: threefry, philox, rbg, unsafe_rbg
• Advanced 10_random_sharding › Exercises
• Advanced 10_random_sharding › 4.1 Default value
• Basics 04_vmap › 4.6 Combining vmap with jit and grad
• Basics 01_introduction › 1.7 Hands-on: basic array operations
• Advanced 10_random_sharding › 4.5 "Random numbers aren't guaranteed to match across versions", and the 10.81 from Chapter 1
• Advanced 10_random_sharding › 5.3 fold_in(key, step): deriving keys for training steps / data indices

State and optimizers

• Advanced 06_large_scale_training › 6.9 Optimization algorithms at scale
• Basics 06_state › 6.6 Example: simple optimizer state
• Advanced 06_large_scale_training › 6.3 Managing model parameters and state
• Basics 06_state › 🔍 Extra: how ecosystem libraries manage state: Optax and Flax
• Basics 07_practical_topics › 5.2 How it works: GradientTransformation and its components
• Advanced 06_large_scale_training › 6.5 Gradient accumulation
• Advanced 11_scientific_computing › 5.2 Optimizing with Optax's L-BFGS
• Basics 07_practical_topics › ★.2 Model, TrainState, and optimizer

Gradient checkpointing (remat)

• Advanced 06_large_scale_training › 6.6 Gradient checkpointing (rematerialization)
• Advanced 06_large_scale_training › 6.10 Practice: implementing gradient checkpointing
• Advanced 06_large_scale_training › 6.3 Managing model parameters and state
• Advanced 04_advanced_autodiff › 4.8 Differentiating through control-flow primitives
• Advanced 11_scientific_computing › 6.3 Memory for long time integration: checkpoint and $O(\sqrt{T})$ segmentation
• Advanced 01_transformations_control_flow › 1.5 Combining control flow with transformations
• Advanced 09_export_deploy › §3.4 Working with Flax NNX: save, restore, keep training
• Advanced 06_large_scale_training › 🧭 Where to go next

Gradient accumulation

• Advanced 06_large_scale_training › 6.5 Gradient accumulation
• Advanced 06_large_scale_training › 6.8 Model parallelism strategies
• Advanced 03_distributed › 3.1 Introduction to parallelism
• Advanced 06_large_scale_training › ✍️ Hands-on exercises
• Advanced 06_large_scale_training › 6.3 Managing model parameters and state
• Advanced 06_large_scale_training › 6.9 Optimization algorithms at scale

Mixed precision

• Advanced 06_large_scale_training › 6.7 Mixed precision training
• Advanced 07_pallas › 4.1′ A scratch accumulator: bf16 inputs, float32 accumulation
• Advanced 06_large_scale_training › 6.1 Challenges in training large models
• Advanced 06_large_scale_training › ✍️ Hands-on exercises
• Advanced 06_large_scale_training › 6.9 Optimization algorithms at scale
• Advanced 07_pallas › 6.2 compiler_params: backend-specific parameters
• Advanced 07_pallas › 6.3 Tuning: block size, arithmetic intensity and the VMEM budget
• Advanced 03_distributed › 3.1 Introduction to parallelism

Memory and donation

• Basics 06_state › 6.6 Example: simple optimizer state
• Advanced 08_debugging › 7.1 memory_analysis(): know how much memory you need before running
• Basics 04_vmap › 4.7 Performance considerations for vmap
• Basics 07_practical_topics › 6.2 How it works: input-output aliasing
• Advanced 02_performance › 2.4 Memory layout and its effect on performance
• Advanced 06_large_scale_training › 6.1 Challenges in training large models
• Advanced 06_large_scale_training › 6.6 Gradient checkpointing (rematerialization)
• Basics 02_jit › 2.6 Common jit pitfalls

Profiling and fusion

• Advanced 02_performance › 2.3 The role of XLA compilation
• Advanced 02_performance › 2.6 Operator fusion and op optimization
• Advanced 02_performance › 2.1 Profiling JAX code on CPU, GPU and TPU
• Advanced 02_performance › 2.7 Asynchronous dispatch
• Advanced 07_pallas › 1.3 Comparing three routes to "custom operations"
• Advanced 09_export_deploy › §4.1 StableHLO text / bytecode → C++ runtimes (IREE, PJRT)
• Basics 02_jit › 2.1 Speeding up: why compile?
• Advanced 02_performance › Hands-on exercises

Debugging tools

• Advanced 08_debugging › 3.2 jax.experimental.checkify
• Basics 05_pmap › 5.7 Debugging pmapped functions
• Basics 02_jit › 2.6 Common jit pitfalls
• Advanced 08_debugging › 1.4 jax.disable_jit(): a debugging aid and a bisection tool
• Advanced 08_debugging › 1.1 jax.debug.print
• Advanced 08_debugging › 1.3 jax.debug.breakpoint
• Basics 04_vmap › 4.2 Introducing jax.vmap
• Advanced 05_interop_custom_ops › 5.3 Calling external CPU/GPU code with jax.experimental.host_callback

Callbacks and external code

• Advanced 05_interop_custom_ops › 5.3 Calling external CPU/GPU code with jax.experimental.host_callback
• Advanced 05_interop_custom_ops › 5.4 Side-effect-free calls with jax.pure_callback
• Advanced 05_interop_custom_ops › 5.10 Practice: integrating a C++ function
• Advanced 08_debugging › 1.2 jax.debug.callback
• Advanced 05_interop_custom_ops › Links to the Basics part
• Advanced 09_export_deploy › §4.4 Wrapping a model as a service function with pure_callback: the trade-offs
• Advanced 05_interop_custom_ops › Hands-on exercises
• Advanced 08_debugging › 1.1 jax.debug.print

Custom primitives and FFI

• Advanced 05_interop_custom_ops › 5.5 Introduction to JAX primitives
• Advanced 05_interop_custom_ops › 5.10 Practice: integrating a C++ function
• Advanced 05_interop_custom_ops › 5.6 Defining custom primitives
• Advanced 02_performance › 2.2 Understanding the JAX computation graph (jaxpr)
• Advanced 05_interop_custom_ops › Hands-on exercises
• Advanced 05_interop_custom_ops › Links to the Basics part
• Advanced 04_advanced_autodiff › 4.8 Differentiating through control-flow primitives
• Advanced 07_pallas › 1.3 Comparing three routes to "custom operations"

DLPack and zero-copy

• Advanced 05_interop_custom_ops › 5.2 Zero-copy data sharing with DLPack
• Advanced 05_interop_custom_ops › Links to the Basics part
• Advanced 05_interop_custom_ops › 5.1 Integrating JAX with NumPy
• Advanced 02_performance › 2.4 Memory layout and its effect on performance
• Basics 07_practical_topics › 6.5 ❌ Common mistakes, measured
• Basics 07_practical_topics › 💡 Practical rule

Pallas kernels

• Advanced 07_pallas › 1.3 Comparing three routes to "custom operations"
• Advanced 07_pallas › 5.3 grad: not differentiable by default; use custom_vjp + a Pallas backward kernel
• Advanced 07_pallas › How this connects to earlier chapters
• Advanced 07_pallas › §0 Probes: confirming the Pallas API in 0.12
• Advanced 07_pallas › 2.2 grid, BlockSpec and pl.program_id
• Advanced 07_pallas › 6.3 Tuning: block size, arithmetic intensity and the VMEM budget
• Advanced 09_export_deploy › §2.6 Version compatibility: three layers of version numbers
• Advanced 07_pallas › 2.1 Kernels, Refs and out_shape

Export and deployment

• Advanced 09_export_deploy › §2.6 Version compatibility: three layers of version numbers
• Advanced 09_export_deploy › §2.5 Exporting differentiable functions: vjp_order
• Advanced 09_export_deploy › §5 Putting it together: train → checkpoint → export → standalone inference in a brand-new subprocess
• Advanced 09_export_deploy › §2.1 Basic usage: what's inside an Exported object
• Advanced 09_export_deploy › §2.3 Shape polymorphism: export once, handle any batch size
• Advanced 09_export_deploy › §2.7 Exporting functions with sharding, shard_map or pmap
• Advanced 09_export_deploy › §4.4 Wrapping a model as a service function with pure_callback: the trade-offs
• Advanced 09_export_deploy › §0 Environment notes and shared helpers

Flax NNX

• Advanced 06_large_scale_training › 6.2 Introduction to JAX ecosystem libraries (Flax, Haiku)
• Basics 06_state › 🔍 Extra: how ecosystem libraries manage state: Optax and Flax
• Advanced 06_large_scale_training › 6.3 Managing model parameters and state
• Advanced 09_export_deploy › §3.4 Working with Flax NNX: save, restore, keep training
• Advanced 06_large_scale_training › 6.4 Combining pmap with training frameworks
• Advanced 10_random_sharding › 5.7 Key management for dropout, and the design of Flax nnx.Rngs
• Advanced 06_large_scale_training › 6.6 Gradient checkpointing (rematerialization)
• Advanced 06_large_scale_training › 6.8 Model parallelism strategies

Scientific computing and PDEs

• Advanced 11_scientific_computing › 3.1 Assembling the five-point matrix from triplets and cross-checking against SciPy CSR
• Advanced 11_scientific_computing › 3.2 sparse.sparsify: letting "dense-style" functions accept sparse matrices
• Advanced 11_scientific_computing › 3.3 BCOO + CG, and the state of sparse direct solvers
• Advanced 04_advanced_autodiff › 4.5 Computing full Jacobians and Hessians
• Advanced 11_scientific_computing › 2.2 The matrix-free operator
• Advanced 11_scientific_computing › 6.1 Explicit and Crank–Nicolson schemes
• Advanced 11_scientific_computing › 6.3 Memory for long time integration: checkpoint and $O(\sqrt{T})$ segmentation
• Advanced 11_scientific_computing › 7.3 Two smoothing energies

Known issues in this environment

• Advanced 02_performance › 2.3 The role of XLA compilation
• Advanced 08_debugging › 0 Symptom lookup table
• Advanced 08_debugging › 3.1 jax_debug_nans / jax_debug_infs
• Basics 03_grad › 3.3 How autodiff works: reverse mode
• Basics 07_practical_topics › 0 Common setup
• Basics 07_practical_topics › 2.4 Locating NaNs with jax_debug_nans (measured behavior in 0.12)
• Advanced 01_transformations_control_flow › How this connects to the Basics part
• Advanced 01_transformations_control_flow › 1.7 Understanding closures and JAX staging

APIs

Grouped by namespace; each API lists the notebooks that cover it (up to 6, linking to the first section where it appears).

jax.*

lax.*

jnp.*

nnx.*

optax.*

ocp.*

pl.*

checkify.*

sparse.*