The Need for Speed: Why Compile?

JAX provides a familiar jax.numpy interface, allowing you to write numerical code that looks much like standard NumPy. However, if you run this code directly, you might notice something underwhelming: it's often not much faster than NumPy, and sometimes it can even be slower. Why is this the case, especially when JAX is touted for high performance on accelerators like GPUs and TPUs?

The answer lies in how standard Python code executes. Python is an interpreted language. When you run a Python function, the interpreter reads, interprets, and executes your code line by line. This dynamic nature offers great flexibility but comes with significant overhead, particularly for numerical computations involving loops or repeated operations on large arrays.

Consider a typical numerical task, like applying a function element-wise to a large array. In pure Python, this would likely involve a for loop. Each iteration of the loop involves interpreter overhead: checking types, looking up methods, handling potential errors, and executing Python bytecode. While NumPy significantly improves this by using pre-compiled C or Fortran routines for array operations, the Python interpreter still plays a role in orchestrating these operations.

Modern hardware accelerators, such as GPUs and TPUs, thrive on executing large, parallel computations. They have thousands of cores designed to perform mathematical operations simultaneously. However, the line-by-line execution model of the Python interpreter becomes a bottleneck. Sending individual instructions or small operations from the CPU to the GPU/TPU incurs latency. The accelerator spends more time waiting for instructions than actually computing. It's like trying to direct a massive construction crew by shouting instructions one worker at a time through a megaphone; the crew is capable of much more if given a complete blueprint beforehand.

This is where compilation becomes essential. Instead of interpreting Python code step-by-step, we can translate a whole function, or a significant part of it, into a lower-level representation (like XLA HLO) optimized specifically for the target hardware (CPU, GPU, or TPU). This compilation process happens Just-In-Time (JIT) when the function is first called with specific input types and shapes.

The compiled function:

1. Reduces Interpreter Overhead: The Python interpreter is largely taken out of the loop during the execution of the compiled code block.
2. Optimizes Operations: The compiler can fuse multiple operations together, eliminating the need to store intermediate results in memory. For example, element-wise operations like $y = a * x + b$ can often be fused into a single kernel launch on the GPU/TPU.
3. Uses Hardware Parallelism: The compiler understands the target architecture (CPU, GPU, TPU) and generates code that takes full advantage of its parallelism. It can schedule large chunks of computation to run efficiently on the accelerator.

The following chart illustrates the potential performance difference between different execution approaches for a numerical task.

Relative speed increase for a typical numerical task when moving from interpreted Python to compiled JAX code. Note the logarithmic scale on the y-axis. Exact speedups vary greatly depending on the task and hardware.

Without compilation, JAX operations executed directly often behave similarly to NumPy operations, relying on dispatching pre-compiled kernels for individual operations but still incurring Python overhead between them. While JAX's potential for speed comes from its ability to compile, this potential is only realized when explicitly requested.

Therefore, to use the performance benefits JAX offers, particularly on accelerators, we need to instruct JAX to compile our functions. The primary tool for this is jax.jit, which we will discuss in detail next. It acts as the "blueprint generator," translating your Python functions into highly optimized code ready for execution.