How JAX Compiles Your Code: The Secret Relationship Between JAX and XLA

In the previous post, we explored why JAX thinks differently, why it separates parameters from computation, and how its functional mindset gives you more control than object-oriented frameworks. But there is still one major question left unanswered.

If JAX is just NumPy with magical transformations, how does it suddenly become so fast on GPUs and TPUs?

To understand that, we need to uncover how JAX compiles your Python code. And the moment you see this clearly, the entire design of JAX makes perfect sense.

Let’s break this down with simple analogies, clear intuition, and real code.

The Big Analogy: JAX as an Architect, XLA as the Construction Team

Imagine you design a building on paper. You draw the blueprint. That blueprint is your Python function.

But the blueprint alone cannot build a skyscraper. You need a construction team that knows how to pour concrete, lift steel, and assemble everything using heavy machinery.

In JAX, the architect is JAX itself. The construction team is XLA.

You write the instructions at a high level. JAX analyzes your blueprint. XLA takes that blueprint and builds the fastest possible version of it for your hardware.

This partnership is the reason JAX feels different from every other framework.

What Actually Happens When You Use jit

Let’s take a simple function.

import jax
import jax.numpy as jnp

def compute(x):
    return jnp.sin(x) + jnp.cos(x)

If you call the function normally, JAX executes it operation by operation.

But when you wrap it with jit:

fast_compute = jax.jit(compute)

JAX does something entirely different. It stops executing your code. Instead, it begins tracing your function.

Tracing is like JAX walking through your function slowly, collecting the mathematical steps you wrote. It does not run them. It records them.

Once JAX collects those steps, it hands them over to XLA, and XLA starts building.

It fuses operations together. It rearranges them for maximum parallelism. It removes unnecessary steps. It compiles everything into one optimized accelerator program.

Then, when you finally call:

y = fast_compute(jnp.ones(1000000))

You are not running Python anymore. You are running pure, optimized machine code on the GPU or TPU.

Why GPUs Love XLA’s Style of Execution

GPUs do not like receiving small, separate tasks. They want one big package of work to run in parallel.

In many frameworks, every operation becomes a separate GPU call. This causes overhead. It is like asking a construction team to build your house brick by brick instead of giving them full walls.

XLA avoids this problem entirely. It merges your operations into a single fused kernel.

So instead of sending: Compute sin Compute cos Add sin and cos

XLA sends: Compute sin(x) + cos(x) in one fused operation This fusion is the secret behind JAX’s speed.

TPU Compilation: Why JAX Fits TPUs Naturally

TPUs cannot interpret Python at all. Everything must be compiled into a TPU-compatible program before execution.

This is where JAX and XLA shine. Since JAX functions are pure and stateless, they are perfectly suited for compilation.

A JAX function that works on CPU will work on GPU. The same function will work on TPU. Nothing needs to be rewritten.

Let’s test this idea in code.

def forward(x):
    return jnp.tanh(x * 3.0)

compiled = jax.jit(forward)

Whether forward runs on CPU, GPU, or TPU depends only on the device. The function itself never changes.

This is why Google researchers often use JAX on TPUs. The compilation pipeline is stable, predictable, and incredibly efficient.

Device Placement: Moving Work to a GPU or TPU

To explicitly send work to a GPU, JAX makes it easy.

gpu = jax.devices("gpu")[0]

def func(x):
    return jnp.sqrt(x)
compiled = jax.jit(func)
x = jnp.ones((1000, 1000))
y = compiled(x).block_until_ready()

If your device is GPU-enabled, the compiled version automatically runs there. JAX selects the most powerful device unless told otherwise.

You can confirm the device by printing:

print(y.device())

It will show something like:

GpuDevice(id=0)

This confirms that the function was compiled for and executed on the GPU.

Why This Compilation Approach Makes JAX Unique

Most frameworks interpret operations eagerly. JAX compiles entire functions. Most frameworks keep internal state hidden inside objects. JAX keeps everything explicit and pure. Most frameworks send many small ops to accelerators. XLA fuses them into large kernels.

This is why JAX feels different. This is why JAX feels fast. This is why JAX scales to extremely large models with fewer surprises.

The functional mindset is not just a programming style. It is the key that unlocks compilation. And compilation is the key that unlocks speed.

What’s Next

In the next post, we will explore **how Flax modules work internally**, why they look object-oriented even though JAX is functional, how the compact API actually works, and how Flax manages parameters under the hood without breaking the functional model.