What jit Does: Tracing, Compiling, Caching

jax.jit is the transformation that makes JAX code fast. It takes a Python function, traces it to extract the computation graph, compiles that graph to optimized machine code using the XLA compiler, and caches the compiled code for reuse.

import jax
import jax.numpy as jnp

def slow_fn(x):
    """Each operation launches a separate kernel."""
    y = jnp.sin(x)
    z = jnp.cos(x)
    return jnp.sum(y * z + y ** 2)

# Wrap with jit
fast_fn = jax.jit(slow_fn)

# Or use as a decorator
@jax.jit
def fast_fn_v2(x):
    y = jnp.sin(x)
    z = jnp.cos(x)
    return jnp.sum(y * z + y ** 2)

x = jnp.ones(10000)
result = fast_fn(x)  # First call: trace + compile + execute
result = fast_fn(x)  # Second call: execute cached compiled code (much faster)

The Three Phases of a JIT Call

1. Tracing: JAX runs your Python function, but instead of using real values, it uses abstract placeholders called tracers. These tracers have shapes and dtypes but no concrete values. JAX records every operation performed on these tracers.
2. Compilation: The recorded operations are sent to the XLA compiler, which optimizes them (fusing kernels, optimizing memory, generating hardware-specific code) and produces a compiled executable.
3. Execution: The compiled code runs on the actual data. This is the fast part.

Phases 1 and 2 only happen on the first call with a given set of input shapes and dtypes (the "signature"). Subsequent calls with matching signatures skip directly to phase 3.

import jax
import jax.numpy as jnp

@jax.jit
def add_one(x):
    print("TRACING!")  # This print helps us see when tracing happens
    return x + 1

# Call 1: traces + compiles + executes
add_one(jnp.array([1.0, 2.0]))  # Prints "TRACING!"

# Call 2: same shape → uses cache, no retracing
add_one(jnp.array([3.0, 4.0]))  # No print — cached!

# Call 3: different shape → must retrace
add_one(jnp.array([1.0, 2.0, 3.0]))  # Prints "TRACING!" — new shape