JIT Compilation for Speedups

1. Eager vs. Compiled Execution

By default, JAX executes operations eagerly, much like NumPy or PyTorch in its default mode. When you execute c = a + b, JAX dispatches the addition operation to the accelerator (CPU, GPU, or TPU) and returns the result. If the next line is d = c * 2, that's a separate, second dispatch. This is simple and intuitive for interactive development, but it's often not the most performant approach. The Python interpreter itself becomes a bottleneck, and the accelerator can't perform optimizations across multiple operations.

The alternative is compiled execution. With jax.jit, you instruct JAX to take an entire Python function, translate it into an optimized, intermediate representation, and then compile it into a single, highly efficient kernel using the XLA (Accelerated Linear Algebra) compiler. This compiled kernel can then be executed with a single dispatch, bypassing Python overhead and enabling powerful optimizations like operator fusion, where multiple operations (e.g., a convolution, bias add, and ReLU activation) are fused into a single, more efficient computation.

2. The JIT Compilation Pipeline: From Python to Machine Code

So what actually happens when you apply @jax.jit to a function? The process involves two distinct paths: a "slow path" for the first execution and a "fast path" for all subsequent executions with similar inputs.

First Call (Slow Path): Tracing and Compilation

When you call a JIT-compiled function for the first time, JAX doesn't just run your Python code on the input data. Instead, it performs a process called tracing:

1. JAX replaces the input arrays with special tracer objects. These are placeholders that have the same shape and dtype as your inputs but no concrete values.
2. It then "traces" the execution of your Python function, with these tracers as inputs. Every JAX operation (jnp.dot, jnp.tanh, etc.) performed on the tracers is recorded.
3. This recording builds a language-agnostic, explicit computation graph known as a jaxpr (JAX expression). This jaxpr captures the pure mathematical computation, stripped of any Python-specific details. Because of this process, side effects like print() or modifying global variables are not captured in the jaxpr and will only execute during this initial tracing step.
4. The jaxpr is then lowered to a portable intermediate representation called StableHLO.
5. This is handed off to the XLA compiler, which performs sophisticated optimizations like kernel fusion, memory layout optimization, and scheduling.
6. XLA compiles the optimized graph into highly efficient machine code for the specific target accelerator (e.g., CUDA for NVIDIA GPUs or TPU code for Google TPUs).
7. Finally, this compiled machine code is stored in a cache, keyed by the function's identity and the signature of the input arguments (i.e., their shapes and dtypes).

This entire pipeline is visualized beautifully in the diagram below.

This diagram shows the full JAX compilation pipeline. The "First Call (Slow Path)" follows the solid arrows from the Python function through tracing, various intermediate representations (IRs), XLA compilation, and finally execution. The resulting machine code is stored in a cache. The "Fast Path" for subsequent calls is the dashed green line, which directly uses the cached code.

Subsequent Calls (Fast Path): The Cache Hit

When you call the same JIT-compiled function again with inputs that have the exact same signature (i.e., same number of arguments with identical shapes and dtypes), JAX simply retrieves the corresponding machine code from its cache and executes it directly on the device. This "fast path" completely bypasses Python, tracing, and XLA compilation, resulting in a dramatic speedup.

This tracing mechanism is why functional purity is so critical. If a function had side effects, the traced jaxpr would not be a faithful representation of what the function does, leading to incorrect behavior.

3. Applying jax.jit and Measuring Performance

Let's put this into practice. We'll define a function, compile it with jax.jit, and measure the performance difference.

A common and correct way to use jit is as a function decorator:

import jax
import jax.numpy as jnp

@jax.jit
def my_jitted_function(x, y):
  # A series of JAX operations
  return jnp.dot(jnp.tanh(x), jnp.sin(y))

Now, let's benchmark. A naive timing approach using time.time() will be misleading due to JAX's asynchronous dispatch. JAX queues operations on the accelerator and returns control to Python immediately, without waiting for the computation to finish. To get an accurate measurement, you must explicitly block the host program until the computation is complete using .block_until_ready() on an output array.

Following this guidance, here is a proper microbenchmark. We'll use a Scaled Exponential Linear Unit (SELU) function as our example, similar to the one in the documentation.

import jax
import jax.numpy as jnp
from jax import random

# 1. Define the function
def selu(x, alpha=1.67, lmbda=1.05):
  return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)

# 2. Create a JIT-compiled version
selu_jit = jax.jit(selu)

# 3. Create some large input data
key = random.key(0)
x = random.normal(key, (5000, 5000))

# --- Eager Execution Benchmark ---
print("Eager execution:")
%timeit selu(x).block_until_ready()

# --- JIT Compilation Cost (First Call) ---
print("\nJIT compilation (first call):")
%time selu_jit(x).block_until_ready()

# --- JIT Cached Execution Benchmark ---
print("\nJIT cached execution:")
%timeit selu_jit(x).block_until_ready()

When you run this (ideally in a Colab/Jupyter environment with a GPU), you will observe:

1. Eager execution: A baseline time, likely in the milliseconds.
2. JIT compilation: A significantly longer time for the first call. This is the one-time cost of tracing, compiling, and caching.
3. JIT cached execution: A dramatically faster time, often an order of magnitude or more, compared to the eager version. This is the performance you'll get for all subsequent calls.

Conclusion

In this lesson, we've taken our first step into JAX's powerful transformation system. You've learned how jax.jit leverages JAX's functional design to provide dramatic performance improvements.

Here are the key takeaways:

• JIT compiles pure functions: jax.jit transforms a Python function into an optimized, compiled kernel via XLA.
• Tracing is the mechanism: On the first call, JAX traces your function with placeholder objects to create a jaxpr, which is then compiled.
• Caching is the payoff: The expensive compilation happens only once per input signature. Subsequent calls are extremely fast as they use the cached machine code.
• Benchmarking requires care: Due to asynchronous dispatch, you must use .block_until_ready() on an output array to get accurate performance measurements.

We've seen the power of jit, but we've also seen that the compilation process is tied to the signature of the inputs. What happens if you call a jitted function with different arguments? Sometimes, JAX must recompile, which can negate the performance benefits. In our next lesson, we will explore this crucial topic, digging into the distinction between trace-time and run-time, what causes re-compilation, and how to control it using tools like static_argnums.