NumPy on Accelerators and XLA Compilation

The phrase "NumPy on accelerators" is more than a tagline — it's JAX's foundational design decision. The jax.numpy module provides an API that mirrors NumPy almost exactly, but every operation can run on GPUs and TPUs, not just CPUs.

import numpy as np
import jax.numpy as jnp

# NumPy — runs on CPU
a_np = np.array([1.0, 2.0, 3.0])
result_np = np.dot(a_np, a_np)  # 14.0

# JAX — same syntax, but can run on GPU/TPU
a_jax = jnp.array([1.0, 2.0, 3.0])
result_jax = jnp.dot(a_jax, a_jax)  # 14.0 (on whatever accelerator is available)

This means if you already know NumPy, you already know most of JAX's computation API. The vast majority of NumPy functions have a jnp equivalent with the same name and the same behavior.

XLA: The Compiler That Makes JAX Fast

Behind the scenes, JAX doesn't execute operations one at a time like NumPy does. Instead, when you use jax.jit, JAX traces your function to build a computation graph and hands it to the XLA (Accelerated Linear Algebra) compiler. XLA is the same compiler that powers TensorFlow's compiled mode, and it's extremely good at optimizing numerical code.

XLA performs several critical optimizations:

• Kernel fusion: Instead of launching a separate GPU kernel for each operation (add, multiply, reduce), XLA fuses multiple operations into a single kernel. This dramatically reduces memory traffic and kernel launch overhead.
• Memory layout optimization: XLA can rearrange how arrays are stored in memory for maximum throughput on the target hardware.
• Whole-program optimization: Because XLA sees the entire computation at once (not just one operation), it can make global decisions about scheduling, memory allocation, and parallelism.
• Hardware-specific code generation: XLA generates optimized code for the specific hardware you're running on — NVIDIA GPUs, Google TPUs, or CPUs each get tailored machine code.

import jax
import jax.numpy as jnp

def computation(x):
    # XLA will fuse these into fewer GPU kernels
    y = jnp.sin(x)
    z = jnp.cos(x)
    return jnp.sum(y * z + y ** 2)

# Without JIT: each operation launches a separate GPU kernel
result = computation(jnp.ones(1000))

# With JIT: XLA fuses operations, dramatically faster
fast_computation = jax.jit(computation)
result = fast_computation(jnp.ones(1000))

The trade-off is compilation time. The first time you call a JIT-compiled function with a new input shape, JAX has to trace and compile it. This can take seconds or even minutes for complex functions. But subsequent calls with the same shapes use the cached compiled code and are extremely fast. This makes JIT a poor fit for one-off computations but perfect for training loops where you call the same function thousands of times.