TopolearnCorrect JAX code can still be slow; these habits often give large speedups for free.
JAX's performance model rests on a few facts.
Asynchronous dispatch. Python enqueues work and returns futures. Anything that needs a value on the host, such as float(x), print(x), np.asarray(x) or a Python if on a result, blocks until the device catches up. Keep metrics on device and fetch them every N steps.
Buffer donation. jax.jit(step, donate_argnums=(0,)) tells XLA it may reuse the input buffers for the outputs. For a train step that maps old params to new params this roughly halves peak parameter memory. Donated inputs must not be used again afterwards.
Compilation cost. Every new shape compiles. Reduce variety (pad to buckets), and enable the persistent compilation cache so restarts reuse executables:
jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")Precision. On GPUs and TPUs matrix multiplies may use reduced-precision passes by default; jax.default_matmul_precision("highest") or the precision= argument trade speed for accuracy. Use bfloat16 activations for speed where accuracy allows.
Profiling. jax.profiler.trace("/tmp/trace") records a trace you can open in Perfetto or TensorBoard to see kernels, gaps and host-device transfers.
Also measure correctly: warm up to exclude compilation, then time with .block_until_ready().
Example
import time
import jax
import jax.numpy as jnp
def step(params, x):
return jax.tree.map(lambda p: p - 0.01 * jnp.tanh(p * x.mean()), params)
fast_step = jax.jit(step, donate_argnums=(0,)) # reuse params' buffers
params = {"w": jnp.ones((2048, 2048))}
x = jnp.ones(16)
params = fast_step(params, x) # compile once
jax.block_until_ready(params)
t = time.perf_counter()
for _ in range(20):
params = fast_step(params, x) # old params are donated
jax.block_until_ready(params) # wait before stopping timer
print(f"{(time.perf_counter() - t) / 20 * 1e3:.2f} ms/step")
lowered = jax.jit(step).lower(params, x)
print(type(lowered.compile()).__name__) # Compiled: AOT compile, no runRemember
- Host reads like float(x) force synchronization; keep values on device
donate_argnumslets XLA reuse input buffers for outputs- Avoid shape variety and enable the persistent compilation cache
Watch out for
- Reading a donated array after the call, which raises an error because its buffer was deleted.