JAX: tracing and transformations
JAX treats differentiation, batching and compilation as transformations of pure functions. See how it traces your Python into a jaxpr with abstract values, transforms that program, and hands it to the XLA compiler — and why that makes side effects, control flow, randomness and mutation work differently.
50 min
After this lesson you can
- Use grad, jit and vmap as composable function transformations
- Explain tracing with abstract values (ShapedArray) and read a jaxpr
- Describe how grad, vmap and jit act on a jaxpr, and what StableHLO and XLA are
- Predict what happens to Python side effects and value-dependent control flow under jit, and use lax.cond and lax.scan
- Use explicit PRNG keys, pytrees and .at[] updates correctly
In this lesson
- Tracing: running Python with stand-in values
- Transformations are interpreters
- The price: pure functions
- Pytrees: structured parameters
- Scaling out: sharding
This lesson is part of Descent Pro
Descent Pro unlocks every lesson on the route, from base camp to the summit. Every plan starts with a free trial.