Need to make a numerical function faster without rewriting it for every device? Just-in-time (JIT) compilation lets JAX compile that function at runtime for a CPU, GPU, or TPU. The first call performs tracing and compilation; later calls reuse optimized code. section 5.1.1 section 5.5
Use JIT in JAX by following this pattern:
- Write a pure function whose behavior depends only on its inputs; avoid global state and side effects. section 5.1.2
- Transform it with jax.jit(function) or decorate it with @jax.jit. section 5.1.1
- Warm up the transformed function before benchmarking, because JAX traces it, creates Jaxpr, compiles with XLA, and caches the result on the first call. section 5.1.2
- Compile a large outer computation when possible, giving the compiler more operations to optimize, fuse, or eliminate. section 1.1.1 section 5.1.1
Compiling a training step
from jax import jit @jit def update(params, x, y, epoch_number): loss_value, grads = value_and_grad(loss)(params, x, y) lr = INIT_LR * DECAY_RATE ** (epoch_number / DECAY_STEPS) return [(w - lr * dw, b - lr * db) for (w, b), (dw, db) in zip(params, grads)], loss_value Decorating the resource-heavy update and evaluation functions preserved model quality while reducing an epoch from about 32.6 seconds to 1.7 seconds in the demonstrated run. The initial iteration remained slower because it included compilation. section 2.7 section 2.2
Account for the main constraints:
- Expect recompilation when static arguments or relevant input types and shapes change; mark non-array values with static_argnums when needed. section 5.1.1 section 5.5
- Do not rely on prints, mutable global state, or other side effects: tracing records the computation, not those effects. section 5.1.2 section 5.3.1
- Measure before adopting JIT for small functions, since compilation, MLIR/XLA overhead, or device data transfer can outweigh the gain. section 5.3.6
- Use helper functions or static parameters for class methods; directly decorating a method can make the instance look like an unsupported abstract array. section 5.3.5
Figure 1.2 Time measurements for different ways of calculating our function f(x). Line 1 is a pure NumPy implementation; line 2 uses JAX on CPU; line 3 is a JAX JIT-compiled CPU version; line 4 uses JAX on GPU; and line 5 is a JAX JIT-compiled GPU version.
Compare pure NumPy, uncompiled JAX, and JIT-compiled JAX across CPU and GPU implementations to see where compilation changes runtime. section 1.1.1
