Topolearnjit is where JAX's speed comes from; uncompiled JAX code dispatches one operation at a time.
jax.jit(f) (or the @jax.jit decorator) compiles f with the XLA compiler. The first call traces the function: JAX runs your Python code with abstract placeholder values that carry only shape and dtype, records the array operations into an intermediate program, and compiles it for the device. Later calls with the same shapes and dtypes skip Python entirely and run the cached executable.
What you gain: operations get fused (no intermediate arrays written to memory), Python overhead disappears, and on GPU/TPU the whole function becomes a few kernels.
What you must accept:
- The function should be pure. Side effects like
printor appending to a list run only during tracing, not on every call. - Values are abstract during tracing, so Python
if x > 0:on an array argument fails with a concretization error. Usejnp.whereorjax.lax.cond(covered later). - A call with a new shape or dtype triggers a fresh trace and compile.
Because dispatch is asynchronous, time jitted code with .block_until_ready(), otherwise you measure only the launch. Always exclude the first (compiling) call from benchmarks.
A good rule: jit the largest pure function you can, typically a whole training or evaluation step, rather than tiny helpers.
Example
import time
import jax
import jax.numpy as jnp
def selu(x, alpha=1.67, lam=1.05):
return lam * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)
selu_jit = jax.jit(selu)
x = jnp.arange(1_000_000, dtype=jnp.float32) / 1e6
selu_jit(x).block_until_ready() # first call: trace + compile
t = time.perf_counter()
selu_jit(x).block_until_ready()
print(f"jit: {time.perf_counter() - t:.5f}s")
t = time.perf_counter()
selu(x).block_until_ready()
print(f"eager: {time.perf_counter() - t:.5f}s") # typically several times slower
@jax.jit
def noisy(x):
print("tracing!") # runs only while tracing
return x * 2
noisy(1.0); noisy(2.0) # prints "tracing!" onceRemember
jax.jittraces once per input shape/dtype and caches compiled XLA code- Jitted functions must be pure; Python side effects run only at trace time
- Benchmark with
.block_until_ready()and skip the compile call
Watch out for
- Jitting tiny functions called from a Python loop, paying dispatch overhead instead of jitting the loop body as a whole.