PyTorch/XLA, From the Inside
You know PyTorch. This is the bridge that runs it on a TPU: every op on the xla device is recorded instead of run, the recorded graph is cut at a sync point, lowered to HLO, and handed to XLA through one C++ class. The track starts where you would, in a Colab notebook, and ends at the exact line where a buffer leaves the host.
Phase 1 — Run it first
Two notebooks on a real TPU, and the craft of writing one that survives a runtime restart.
- LAB·P4 · Torch on a TPU ↗ — Both bridges from torch to a TPU in one runtime each: torchax, then a clean restart, then torch_xla with a training loop, a checkpoint, a kill, and a resume.
- Writing torch_xla Notebooks That Survive Colab — The craft layer: the version pair, what PJRT_DEVICE does when you leave it unset, why two bridges cannot share a runtime, and the three instruments that prove any of it.
- LAB·P6 · The seam trace on a CPU ↗ — The same block traced on the in-process CPU client: lazy IR, HLO, one sync, one read, and the counters that name every call across the seam.
Phase 2 — The lazy tensor
Record instead of run. Then the cut, the hash, the compile cache, and the three other ways to execute the same graph.
- The Lazy Tensor: What Happens Between Your Op and sync() — What a tensor on the xla device actually holds, where the graph gets cut, why the hash decides whether you compile again, and how eager mode and torch.compile ride the same machinery.
- Inside the lazy tensor ↗ — The kernels-site lesson that first names the sync-point model, with the counters to watch.
- Where the graph gets cut ↗ — Every trigger that forces a sync before you asked for one, and what each costs.
- torch_xla/torch_xla.py ↗ — sync() and step(): the Python side of the cut is thirty lines, and every one of them calls into the C++ extension.
- torch_xla/csrc/ir.cpp ↗ — The XlaNode constructors. Interior nodes hash op, seed, and operands; the leaf constructor folds the shape in. That asymmetry is why input shapes reach the compile cache.
Phase 3 — The seam, in theory
One interface, one client per process, and the round trip of a single step across it, with the counters that make it observable.
- The PJRT Boundary: One Training Step, Crossing by Crossing — What crosses the seam (a computation, buffers, a literal), what never does (collectives, donation), and the fourteen files that reach it.
- LAB·P5 · The boundary trace ↗ — One training step traced three ways, lazy, eager, and under torch.compile, with the seam table as the reading frame and the IR and HLO dumps as the records.
Phase 4 — The seam, line by line
The three files that are the seam: how a device string becomes a client, the interface, and the client that calls PJRT.
- How PJRT_DEVICE Becomes a Client: pjrt_registry.cpp, Line by Line — PJRT_DEVICE in, an xla::PjRtClient out: the in-process CPU client, libtpu loaded as a plugin, and any other device through the C API.
- The Seam Interface: computation_client.h, Line by Line — Read the contract before the implementation: 43 pure-virtual methods, the four different things called "Computation" that meet in one class, and the GIL warning on the way back.
- Where torch_xla Calls PJRT: pjrt_computation_client.cpp, Line by Line — The implementation, read end to end: a host buffer in, a compiled executable, one execute per device, and the literal that comes back.
- torch_xla/csrc/runtime/runtime.cpp ↗ — The client is built once per process, and the IFRT alternative is switched off by a hard-coded false, not a flag.
- xla/pjrt/pjrt_client.h ↗ — The other side of the seam at the commit torch_xla pins: Compile returns an executable, CompileAndLoad returns the loaded one torch_xla asks for.
Phase 5 — Above the seam
The executor that collects, hashes, lowers, caches, and runs, and the dynamo bridge that replays a cached graph without tracing.
- From Pending IR to a Device Buffer: xla_graph_executor.cpp, Line by Line — CollectSyncTensors, the six-term graph hash, LookupCachedCompile, Compile, and the async run: the file that turns pending IR into a PJRT call.
- torch.compile Meets the Lazy Tensor: dynamo_bridge.py, Line by Line — The openxla backend does not lower FX nodes. It runs the graph once through the lazy tensor, keeps the hash, and replays by hash on every later call.
- torch_xla/csrc/lowering_context.cpp ↗ — BuildXla: where the post-ordered IR becomes an xla::XlaComputation through an XlaBuilder.
Phase 6 — Many devices
A mesh, a sharding annotation, the virtual SPMD device, and the replicated execute that follows.
- SPMD in torch_xla: A Mesh, an Annotation, and One Virtual Device — Mesh and mark_sharding, the permutation that turns a partition spec into a tile assignment, what the annotation becomes in the HLO, and why every compile is sharded once the virtual device is on.
- Marking a sharding ↗ — The kernels-site lesson that puts a sharding on a tensor and reads the result back from the device.
- The SPMD Partitioner: One Program Across a Device Mesh — The compiler pass on the far side of the seam that turns one annotated program into N per-device programs. torch_xla only sets the options that switch it on.
- Sharding in JAX — The same mesh and partition-spec idea in JAX, useful because torch_xla borrowed the vocabulary and the compiler beneath it.