Benchmarking JAX code requires care because of asynchronous dispatch. JAX operations return immediately and execute asynchronously on the accelerator. To get accurate timing, you must use block_until_ready().
import jax
import jax.numpy as jnp
import time
def matrix_computation(A, B):
"""A non-trivial computation to benchmark."""
C = A @ B
D = jnp.sin(C) + jnp.cos(C)
E = D @ D.T
return jnp.sum(jnp.log(jnp.abs(E) + 1e-7))
# Create test data
key = jax.random.PRNGKey(0)
A = jax.random.normal(key, (1000, 1000))
B = jax.random.normal(jax.random.PRNGKey(1), (1000, 1000))
# Benchmark WITHOUT jit
# Important: block_until_ready() ensures the computation is actually done
def time_fn(fn, *args, n_runs=10):
# Warm up
result = fn(*args)
result.block_until_ready()
start = time.perf_counter()
for _ in range(n_runs):
result = fn(*args)
result.block_until_ready()
elapsed = time.perf_counter() - start
return elapsed / n_runs
eager_time = time_fn(matrix_computation, A, B)
print(f"Eager: {eager_time * 1000:.2f} ms")
# Benchmark WITH jit
jitted_computation = jax.jit(matrix_computation)
# Warm up (trigger compilation)
jitted_computation(A, B).block_until_ready()
jit_time = time_fn(jitted_computation, A, B)
print(f"JIT: {jit_time * 1000:.2f} ms")
print(f"Speedup: {eager_time / jit_time:.1f}x")
Typical results on a GPU might show 3-10x speedup for this kind of fused computation. The speedup comes from XLA fusing the sin, cos, addition, and matmul operations into fewer kernel launches.
Measuring Compilation Time
import jax
import jax.numpy as jnp
import time
@jax.jit
def big_fn(x):
for _ in range(20):
x = jnp.sin(x) + jnp.cos(x)
return jnp.sum(x)
x = jnp.ones((2000, 2000))
# First call includes compilation
start = time.perf_counter()
result = big_fn(x)
result.block_until_ready()
first_call = time.perf_counter() - start
# Second call uses cache
start = time.perf_counter()
result = big_fn(x)
result.block_until_ready()
second_call = time.perf_counter() - start
print(f"First call (compile + run): {first_call * 1000:.1f} ms")
print(f"Second call (run only): {second_call * 1000:.1f} ms")
print(f"Compilation overhead: ~{(first_call - second_call) * 1000:.1f} ms")
💡 Why This Matters
Always use .block_until_ready() when benchmarking JAX. Without it, you're only timing the dispatch of the operation, not the actual computation. This is the most common benchmarking mistake in JAX — people report unrealistically fast times because they measured dispatch overhead instead of actual execution time.
Honest assessment: For simple operations on small arrays, JIT may not help or can even be slower due to dispatch overhead. JIT shines on complex computations with many chained operations, especially on large arrays. The sweet spot is training step functions, which combine forward pass, loss, backward pass, and parameter updates into a single compiled program.