Showing jax.jit on an NVIDIA GPU: 12 stages, 13 steps.
Without jit, Python stays in charge of every operation: each jnp call is dispatched to the device on its own, so the compiler never sees more than one operation at a time.
With jax.jit, Python runs once to trace the function. From then on the computation is a jaxpr, then a StableHLO program, then an executable that XLA built for the device, and later calls with the same shapes reuse the compiled executable without running the Python body again.
Both JAX and PyTorch's torch.compile capture a whole computation and compile it, and both can end up in Triton and LLVM on NVIDIA GPUs. JAX hands its program to XLA through a stable intermediate language, StableHLO, and runs it through PJRT; torch.compile captures an FX graph and generates kernels with TorchInductor. See what happens when you run a PyTorch model.
programming language
Python functions using jax.numpy and other JAX APIs.
In the atlas: Python
Without jit, JAX executes operations (think jnp function calls) one at a time, dispatching each to the device as it's encountered.Just-in-time compilation (JAX documentation) Project docs
framework component
The JAX library. It traces and transforms your functions, and calls out to XLA and PJRT to compile for and run on devices.
To manage hardware devices and memory, and for compilation to such devices, JAX calls out to the open XLA compiler and the PJRT runtime.About the project (JAX documentation) Project docs
intermediate representation
JAX's own intermediate language: the traced function as a list of primitive operations, specialized to the argument shapes and dtypes.
as a jaxpr: a function in JAX's intermediate language.Ahead-of-time lowering and compilation (JAX documentation) Project docs
intermediate representation
XLA's versioned input language, a portability layer between ML frameworks and compilers. It is an operation set, not a programming language.
Lower this specialized, staged-out computation to the XLA compiler's input language, StableHLO.Ahead-of-time lowering and compilation (JAX documentation) Project docs
StableHLO defines a versioned operation set (HLO = high level operations) that provides a portability layer between ML frameworks and the compiler.XLA architecture (OpenXLA) Project docs
compiler
OpenXLA's machine-learning compiler. It optimizes the whole program, fusing operations, and generates code for the target device.
The XLA compiler takes model graphs from ML frameworks defined in StableHLO and compiles them into machine instructions for various architectures.XLA architecture (OpenXLA) Project docs
It traces the function to a jaxpr (How transformations work: tracing) and hands the whole jaxpr to XLA, which compiles it into a single optimized executable, fusing operations, eliminating temporary arrays, and specializing for your CPU, GPU, or TPU.Just-in-time compilation (JAX documentation) Project docs
intermediate representation
XLA's CPU and GPU backends emit LLVM IR and use LLVM to generate native code.
In the atlas: LLVM
The CPU and GPU backends included with XLA use LLVM for low-level IR, optimization, and code generation.XLA architecture (OpenXLA) Project docs
These backends emit the LLVM IR necessary to represent the HLO computation in an efficient manner, and then invoke LLVM to emit native code from this LLVM IR.XLA architecture (OpenXLA) Project docs
intermediate representation
For some GPU kernels, XLA:GPU emits Triton IR instead of going straight to LLVM.
In the atlas: Triton
XLA:GPU uses a combination of "native" (PTX, via LLVM) emitters and TritonIR emitters to generate high-performance GPU kernelsXLA:GPU architecture overview (OpenXLA) Project docs
kernel
NVIDIA's libraries, which XLA:GPU calls for many common operations instead of generating code.
In the atlas: CUDA
For many common operations, XLA:GPU uses fast-performance libraries from NVIDIA, such as cuBLAS, cuDNN, and NCCL.XLA:GPU architecture overview (OpenXLA) Project docs
For every HLO instruction in the computation, XLA chooses whether to run it using a library linked into a runtime, or to codegen it to PTX.XLA:GPU architecture overview (OpenXLA) Project docs
intermediate representation
NVIDIA's virtual instruction set, produced through LLVM's NVPTX backend.
In the atlas: CUDA
The GPU backend currently supports NVIDIA GPUs via the LLVM NVPTX backend.XLA architecture (OpenXLA) Project docs
XLA:GPU uses a combination of "native" (PTX, via LLVM) emitters and TritonIR emitters to generate high-performance GPU kernelsXLA:GPU architecture overview (OpenXLA) Project docs
compiler
XLA compiles for TPUs too. How its TPU backend generates code is not covered by the open documentation cited here.
Compile the lowered HLO program to produce an optimized executable for the target device (CPU, GPU, or TPU).Ahead-of-time lowering and compilation (JAX documentation) Project docs
machine code
The optimized program for one device and one set of argument shapes. Later calls with the same shapes reuse it.
Compile the lowered HLO program to produce an optimized executable for the target device (CPU, GPU, or TPU).Ahead-of-time lowering and compilation (JAX documentation) Project docs
runtime
OpenXLA's uniform device API. Frameworks call PJRT; each device type implements it as a plug-in.
PJRT C API is the uniform Device API that we want to add to the ML ecosystem.PJRT: uniform device API (OpenXLA) Project docs
XLA aims for interoperability across accelerators (e.g. by ingesting StableHLO as input) and PJRT offers extensibility through a plug-in device API.About the project (JAX documentation) Project docs
hardware
XLA's CPU backend supports several CPU instruction sets.
The CPU backend supports multiple CPU ISAs.XLA architecture (OpenXLA) Project docs
hardware
XLA reaches NVIDIA GPUs through LLVM's NVPTX backend and NVIDIA libraries.
The GPU backend currently supports NVIDIA GPUs via the LLVM NVPTX backend.XLA architecture (OpenXLA) Project docs
hardware
Google's Tensor Processing Unit, one of the devices JAX compiles for.
Compile the lowered HLO program to produce an optimized executable for the target device (CPU, GPU, or TPU).Ahead-of-time lowering and compilation (JAX documentation) Project docs
| From | Step | To | When | What happens | Source |
|---|---|---|---|---|---|
| Your JAX code | calls | JAX | any mode · any device | Your functions call JAX's array operations. | Project docs |
| JAX | dispatches each op through | PJRT runtime | Without jit · any device | Without jit, each operation is dispatched to the device on its own, as Python reaches it. | Project docs Project docs |
| JAX | traces to | jaxpr | jax.jit · any device | jax.jit runs the function once with tracers and records it as a jaxpr, specialized to the argument shapes and dtypes. | Project docs Project docs |
| jaxpr | lowers to | StableHLO | jax.jit · any device | The jaxpr is lowered to StableHLO, the XLA compiler's input language. | Project docs |
| StableHLO | compiled by | XLA compiler | jax.jit · any device | XLA compiles the whole StableHLO program, fusing operations and specializing it for the device. | Project docs Project docs |
| XLA compiler | emits | LLVM IR | jax.jit · CPU or NVIDIA GPU | XLA's CPU and GPU backends emit LLVM IR for the computation. | Project docs |
| XLA compiler | emits, for some kernels | Triton IR | jax.jit · NVIDIA GPU | XLA:GPU generates some kernels through Triton IR emitters. | Project docs |
| XLA compiler | calls, for common ops | cuBLAS and cuDNN | jax.jit · NVIDIA GPU | For each operation XLA chooses between generating code and calling a library; many common operations go to cuBLAS or cuDNN. | Project docs Project docs |
| XLA compiler | compiles with | XLA TPU backend | jax.jit · TPU | For TPUs, XLA's TPU backend produces the executable; its internals are not covered here. | Project docs |
| LLVM IR | compiled to native code | Compiled executable | jax.jit · CPU | On CPU, LLVM turns the IR into native machine code. | Project docs Project docs |
| LLVM IR | lowers to | PTX | jax.jit · NVIDIA GPU | On NVIDIA GPUs, LLVM's NVPTX backend turns the IR into PTX. | Project docs Project docs |
| Triton IR | lowers to | PTX | jax.jit · NVIDIA GPU | Triton-emitted kernels also end up as PTX. | Project docs |
| PTX | built into | Compiled executable | jax.jit · NVIDIA GPU | The generated kernels become part of the compiled executable. | Project docs |
| cuBLAS and cuDNN | called from | Compiled executable | jax.jit · NVIDIA GPU | Library calls are linked into the runtime and invoked by the executable. | Project docs |
| XLA TPU backend | produces | Compiled executable | jax.jit · TPU | The TPU backend produces the executable for the TPU. | Project docs |
| Compiled executable | run through | PJRT runtime | jax.jit · any device | JAX runs the compiled executable on the device through the PJRT runtime API. | Project docs Project docs |
| PJRT runtime | runs on | CPU | any mode · CPU | PJRT's CPU implementation runs the work on the host CPU. | Project docs |
| PJRT runtime | launches on | NVIDIA GPU | any mode · NVIDIA GPU | PJRT's GPU implementation launches the work on the NVIDIA GPU. | Project docs Project docs |
| PJRT runtime | launches on | TPU | any mode · TPU | PJRT's TPU implementation runs the work on the TPU. | Project docs Project docs |
Documented paths for one JAX program run without jit and with jax.jit, on a CPU, an NVIDIA GPU, or a TPU. Multi-device sharding, Pallas kernels, AMD GPUs, PJRT plugins for other hardware, and the internals of XLA's TPU backend are out of scope. XLA's TPU backend is not described in the open-source documentation cited here, so this page shows that XLA compiles for TPUs without showing how.
Every stage and step above quotes the source that supports it. The quotes were checked against the JAX documentation (About, Just-in-time compilation, Ahead-of-time lowering) and the OpenXLA documentation (XLA architecture, XLA:GPU architecture, PJRT) on 2026-11-11. The data is a versioned JSON file under the same CC BY 4.0 license as the dataset. Found something wrong or out of date? Open an issue with the better source.
What these tools are written in is in the main atlas: Python, LLVM, Triton, CUDA. More execution paths: How Does llama.cpp Run a Model on Your Hardware, What Actually Happens When You Run a PyTorch Model.
Explore the full graph →