Profiling and Debugging Performance

When your JAX code is slow, you need to figure out why. JAX provides several profiling tools to identify bottlenecks.

Basic Profiling

import jax
import jax.numpy as jnp

# Check what's being compiled
@jax.jit
def my_function(x):
    return jnp.sin(x) + jnp.cos(x)

# Inspect the jaxpr (JAX's intermediate representation)
jaxpr = jax.make_jaxpr(my_function)(jnp.ones(3))
print(jaxpr)
# Shows the operations JAX will compile — useful for checking
# that your function isn't doing unexpected work

# Inspect compiled HLO
compiled = jax.jit(my_function).lower(jnp.ones(3)).compile()
print(compiled.cost_analysis())
# Shows estimated FLOPs, memory usage, etc.

TensorBoard Profiling

import jax

# Profile with TensorBoard integration
jax.profiler.start_trace('/tmp/jax_profile')

# Run your training code
for step in range(100):
    params, loss = train_step(params, batch)

jax.profiler.stop_trace()
# Then: tensorboard --logdir=/tmp/jax_profile

# Or use the context manager
with jax.profiler.trace('/tmp/jax_profile'):
    for step in range(100):
        params, loss = train_step(params, batch)

Detecting Recompilation

# Set this to log when JIT recompiles
jax.config.update("jax_log_compiles", True)

# Now you'll see messages like:
# "Compiling my_function (..." every time JIT compiles a new variant

# Check XLA compilation logs for detailed info
# JAX_LOG_COMPILES=1 python my_script.py

# Timing best practices
import time

@jax.jit
def fast_fn(x):
    return jnp.linalg.svd(x, full_matrices=False)

x = jax.random.normal(jax.random.key(0), (1000, 500))

# First call includes compilation time
start = time.time()
result = fast_fn(x)
jax.block_until_ready(result)
print(f"First call (includes compile): {time.time() - start:.4f}s")

# Second call is the real runtime
start = time.time()
result = fast_fn(x)
jax.block_until_ready(result)
print(f"Second call (actual runtime): {time.time() - start:.4f}s")