Most deep learning engineers meet PyTorch as an eager framework: every op fires a kernel the moment you call it. On a TPU, that model of execution does not work. TPUs are fed by a compiler — XLA, the Accelerated Linear Algebra compiler that also powers JAX and TensorFlow — and a compiler needs to see a whole computation before it can schedule buffers, fuse ops, and tile matrix multiplies across a systolic array. PyTorch/XLA (the torch_xla package, developed in the pytorch/xla repository) is the machinery that makes an eager framework speak fluent compiler.

The deal it offers is simple to state and subtle to implement. You keep writing ordinary PyTorch code — modules, autograd, optimizers, data loaders — but tensors live on a special device called xla, and nothing executes when you call an op. Instead, every op appends a node to a lazily-traced graph. When you call xm.mark_step(), the accumulated graph is lowered to XLA HLO, compiled, cached against a hash, and pushed to the device for execution. This is the same architectural bargain JAX makes explicitly with jit, but grafted onto PyTorch’s eager semantics so that existing code mostly runs unchanged.

This post is a source tour of how that bargain is honored, based on a fresh checkout of pytorch/xla. We will walk the full stack: the xm Python API, the C++ lazy-tensor core that traces and lowers graphs, the PJRT runtime that owns device buffers, the SPMD sharding layer for multi-chip training, and the newer Dynamo bridge that lets torch.compile target XLA directly.

The PyTorch/XLA stack: Python entry points on the left, the C++ lowering core in the middle, the PJRT runtime underneath, and the SPMD sharding layer wrapping it all for multi-device runs.

Reading the overview from left to right:

• xm.xla_device, L122 is the front door: it hands you a torch.device('xla') that any PyTorch model can be moved to.
• xm.mark_step, L1031 is the synchronization beat of the whole system — the point where traced graphs get compiled and shipped.
• set_device_type, L41 in the xr runtime module selects the PJRT device type (CPU, TPU, GPU plugin) and queries world size and ordinals.
• extract_compiled_graph, L670 is the Dynamo bridge that converts a captured FX graph into XLA IR instead of tracing op-by-op.
• SyncTensorsGraphInternal, L1533 is the C++ heart: it hashes the traced graph, hits the compilation cache, or compiles fresh HLO.
• LoweringContext::GetOutputOp, L255 turns lazy IR nodes into concrete XLA ops during lowering.
• PjRtComputationClient::Compile, L544 hands the HLO module to the PJRT client, which executes it against real device buffers.
• mark_sharding, L685 annotates a tensor with an OpSharding spec so the SPMD partitioner can split one program across a mesh of chips.

Lazy Tensors: Pay Only When You Sync

The first thing to internalize is that an xla tensor is not a pointer to device memory the way a CUDA tensor is. In the C++ core, an XLATensor holds either device data, an IR node describing a pending computation, or a view — and the three states migrate lazily. When you compute loss = model(x).sum(), nothing executes. When you call loss.backward() and then xm.optimizer_step(...), the framework accumulates the forward graph, the backward graph, and the gradient update into one big computation, then lowers all of it at the next mark_step.

This design has a beautiful consequence: the unit of compilation is whatever happens between two mark_step calls, and PyTorch/XLA compiles that whole region at once, fusing elementwise chains and materializing only what it must. It also has an obvious cost, which we will return to in the limits section: tracing happens on the host, in C++, one op at a time.

The xm module is the Python-facing contract. Beyond xla_device and mark_step, it carries the collectives that predate the SPMD layer: xm.all_reduce, L396 reduces gradients across replicas with a configurable REDUCE_SUM/REDUCE_MEAN type, and xm.optimizer_step, L1169 bundles gradient all-reduce, the optimizer step, and an optional barrier into one call. The top-level torch_xla package wraps these in a cleaner surface — torch_xla.sync, L78 is the modern spelling of mark_step.

Inside the C++ Lowering Layer

Everything below the Python API happens in torch_xla/csrc. The flow from Python call to device execution is a chain of four stages, each visible in the source.

Stage one is dispatch. InitXlaBackend, L254 registers XLA as a backend with PyTorch’s dispatcher, so calling an op on an xla tensor routes into torch_xla’s implementations rather than ATen kernels. Stage two is tracing: each op implementation in tensor_methods.cpp appends a node to the lazy IR graph — for collectives, tensor_methods.cpp all_reduce, L750 constructs an AllReduce node carrying the reduction type and replica groups, and the node kinds themselves are declared in ops/xla_ops.cpp, L5 as OpKindWrapper constants like xla::cross_replica_sum.

Stage three is the graph executor. When a sync fires, XLAGraphExecutor collects the live tensors for the device, walks the IR in post-order through a LoweringContext, and produces an HLO computation. Crucially, the executor hashes the graph — parameters, IR structure, even auto-sharding configuration — and consults a cache of compiled computations before compiling anything. A cache hit skips straight to execution; this is why steady-state training on XLA is fast even though step one pays a large compile bill. The hash-then-compile-then-schedule sequence lives in TryRunCachedSync, L1245 and the Compile routine it guards.

Stage four is the runtime boundary. The lowering output is an HLO module; the device wants buffers and an executable. The singleton client returned by GetComputationClient, L53 is a PjRtComputationClient (or, experimentally, an IFRT client), and its Compile and TransferToDevice methods are the seam where PyTorch data structures become PJRT buffers. PJRT is the same device API JAX uses, which is what makes cross-framework tooling on TPUs possible at all.

From Install to Your First Model

The README’s stable-build instructions for a TPU VM are refreshingly short:

pip install torch==2.8.0 'torch_xla[tpu]==2.8.0'

The tpu extra pulls in libtpu, the PJRT plugin that talks to the chip. On CPU-only machines (or GPUs via plugins), the package installs from source or from project-specific indexes; there is a full matrix of wheels in the repository README. Once installed, a minimal end-to-end training step uses only APIs you can find in the tree:

import torch
import torch_xla.core.xla_model as xm

device = xm.xla_device()
model = torch.nn.Linear(128, 64).to(device)
opt = torch.optim.SGD(model.parameters(), lr=0.1)

x = torch.randn(16, 128, device=device)
loss = model(x).sum()      # traced, not executed
loss.backward()            # backward graph traced too
xm.optimizer_step(opt)     # all-reduce grads + step
xm.mark_step()             # <- compile & execute happens HERE

The mental model to keep: the loop body is trace-compile-cache-execute. On the first step you pay the compiler; on later steps, the graph hash matches and the cached executable runs. xm.optimizer_step is doing double duty — it wraps reduce_gradients plus optimizer.step() — which is why real examples like test/test_train_mp_mnist.py in the repository call exactly this sequence.

Scaling Out: PJRT Processes and SPMD Meshes

Multi-device execution changed shape a few years ago when torch_xla moved from XRT (an RPC server model) to PJRT. In the PJRT world, replication is process-based: one process per device (or one process hosting several devices), and the Python-level launcher for that is pjrt.spawn, L186, which xmp.spawn in xla_multiprocessing.py, L10 forwards to. Each process binds to its device, and collectives like xm.all_reduce coordinate across processes through the runtime rather than through a central server. Data feeding uses MpDeviceLoader (parallel_loader.py, L253), which shards batches per device and issues the step-closure sync at the right moment. If you have used Accelerate or DeepSpeed, the mental model of “launcher plus wrapped loader plus gradient sync” will feel familiar.

The newer story is SPMD: instead of N processes each running a full program copy, you run one program across a mesh of devices and let the XLA compiler partition it. You opt in with xr.use_spmd(), build a mesh, and annotate tensors. The Mesh class (xla_sharding.py, L38) takes device IDs, a mesh shape, and axis names; get_1d_mesh covers the simple data-parallel case; and mark_sharding attaches an OpSharding to a tensor. The DTensor-flavored entry points are xla_distribute_tensor and xla_distribute_module in spmd/api.py, L110, which accept PyTorch DeviceMesh objects and partition specs — deliberately mirroring the upstream torch.distributed.tensor API. Under the hood the annotations flow into xla_sharding_util.cpp, which propagates sharding through the IR so the HLO SPMD partitioner can slice the computation.

The Dynamo Bridge: torch.compile on XLA

Lazy tracing was built one op at a time, in C++, and it shows in CPU profiles for big models. The modern alternative is to let PyTorch’s own compiler capture the program. torch_xla.compile (torch_xla.py, L114) uses the Dynamo bridge in _dynamo/dynamo_bridge.py: Dynamo produces an FX graph, the bridge runs extract_compiled_graph to convert FX nodes into XLA IR — including an analysis pass that partitions nodes into XLA-executable and CPU-fallback regions — and returns a callable that executes the lowered computation. Because Dynamo’s symbolic reasoning handles dynamic shapes and guards more gracefully than op-by-op tracing, this path is where most new investment is going.

There is also an experimental second backend, dynamo_backend2.py, which routes FX graphs through torchax and JAX’s compiler tooling instead of the native C++ tracer — evidence that the long-term direction is converging on the shared OpenXLA stack that also serves IREE and the rest of the XLA ecosystem, much as TensorFlow standardized on XLA years ago.

Honest Limits

Lazy tensors are a tax as well as a gift. Tracing runs op-by-op on the host, and for large models the host can become the bottleneck — the repository’s own performance docs discuss “tracing-bound” workloads where TPUs sit idle waiting for the CPU, and why the C++11-ABI wheels measurably improve flops utilization on exactly those workloads.

Recompilation is the second recurring cost. The cache is keyed by graph hash, so dynamic shapes (variable sequence lengths, data-dependent control flow) produce fresh hashes and fresh compiles. Code that reshapes itself every step can spend all its time in the compiler; the standard mitigations are padding, bucketing, Dynamo’s shape handling, or explicit graph breaks.

Hardware access is the third practical constraint. TPU code requires a TPU VM (or the CPU/GPU PJRT plugins, which are functional but not the main audience), the wheels are Linux-only — Windows users need WSL — and libtpu is distributed for specific Python versions, so version pinning between Python, torch, and torch_xla matters. Finally, some things you take for granted on CUDA tensors (certain in-place mutations and views) fall back to CPU; the metrics module (torch_xla.debug.metrics) will tell you when a fallback fired, and it is worth checking after your first run.

None of these limits are fatal — Google trains Gemini-scale models through this stack — but they shape a different engineering rhythm than eager GPU work: you profile compilation counts and graph hashes, not just kernel times.

Where PyTorch/XLA Fits

In the widening zoo of execution backends, torch_xla occupies a specific niche: it is the sanctioned path from eager PyTorch semantics onto XLA-compatible accelerators, sharing compiler infrastructure with JAX while keeping your training loop recognizably PyTorch. If your deployment target is a TPU pod, it is not optional; if your target is a laptop or a phone, alternatives like Executorch fit better. But understanding how a lazily-traced graph becomes an HLO program that a compiler partitions across a mesh is one of the most transferable mental models in modern ML infrastructure — and the pytorch/xla source tree, unusually for this kind of project, lets you read every layer of it in one afternoon.

Enjoyed this post? Never miss out on future posts by following us