Performance

Time a JAX op the way you'd time any other Python call, and the number that comes back is very often wrong, not because the clock lied but because of what you handed it to measure. Three mistakes produce three different lies. Time a call without blocking on the result and you measure dispatch, the moment the Python line returns control, not the moment the device finishes computing (chapter 1 built this fact once already: JAX dispatches asynchronously, and printing or converting to NumPy is what actually blocks). Time a function's first call and you measure compilation, the one-time cost of tracing and lowering that every later call skips. Time a function once and you measure noise, whatever the scheduler happened to be doing at that instant.

The harness below kills all three lies at once. A warmup call forces the compile and the first execution to finish before the clock starts. The timed loop runs the function n times and blocks once at the end, so dispatch overhead amortizes across every call instead of dominating a single one. Dividing by n gives a number that reflects the computation, not the ceremony around it.

This site's own habit for every measured number carries over directly: state the machine, the dtype, the shapes, and the method, every time. A latency number with no provenance isn't a fact. It's a rumor with decimal places.

Benchmark the computation, not the dispatch.

import time

import jax

def bench(f, *args, n=20):
    jax.block_until_ready(f(*args))   # warmup: compile + first run
    t0 = time.perf_counter()
    for _ in range(n):
        out = f(*args)
    jax.block_until_ready(out)
    return (time.perf_counter() - t0) / n

A training loop that recompiles mid-run is losing seconds it will never get back, and the first step is knowing it's happening at all. jax.config.update("jax_log_compiles", True) logs every trace and compile event along with the reason JAX decided a new one was needed. Once training reaches steady state, that log should go silent: same shapes, same structure, same executable, call after call.

Chapter 3's cache-key model explains every entry that keeps showing up in that log. Shapes drifting from the data forces a retrace, and the last batch of an epoch is usually short, which makes it the most common offender. A fresh closure or lambda built inside the loop counts as a new function identity even when its body is identical, forcing another retrace. A Python float that flips a leaf from f32 to a weakly typed float across calls changes the key too, quietly, with no exception to flag it.

The standard fix for ragged data, the short last batch above all, is padding to a fixed set of shapes: pad the tail batch up to the usual size before it reaches the traced function, and the cache key stops moving.

Logging compiles tells you when JAX rebuilt something. It doesn't tell you where the time inside a single step actually went. jax.profiler.trace(dir), or the paired start_trace and stop_trace calls, captures a trace you can open in XProf or Perfetto, and that trace shows dispatch gaps, compile events, and the cost of individual ops laid out on a timeline.

jax.profiler.save_device_memory_profile does the same job for memory instead of time, and it's how you go looking for the residuals chapter 4 predicted reverse mode would store. A memory profile that spikes exactly where a vjp pass should be holding its intermediates confirms the cost model rather than just gesturing at it.

The kernel path's IR-stage chapter reads these same profiles down to individual bytes moved. This chapter only needs the workflow: capture a trace, open it, find the gap, not that depth.

The dtype you compute in is a decision with a cost, not a detail the framework hides from you. On TPU, f32 matmuls run at reduced precision on the MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu → by default, trading some accuracy for throughput without asking. When you're comparing a computation against a reference and need the true f32 answer, jax.default_matmul_precision("highest") pins it, the same lesson the kernel path's gate 3 measures directly and publishes at /bench.

Dropping to bf16 halves the bytes moved for every array involved, so memory-bound operations speed up by roughly 2x, and compute-bound matmuls ride the MXUThe systolic matmul array: 128x128 on v5e, 256x256 on v6e. Matmuls only; everything else is the VPU’s job.taught in /l/tpu →'s native bf16 path instead of emulating it. Going the other direction, enabling x64 doubles the bytes and, on accelerators, can push a computation off the fast path entirely rather than just slowing it proportionally.