PyTorch/XLA

The PyTorch frontend for XLA devices: lazy tensors, an HLO lowering, and the PJRT client that runs them on a TPU.

Explainers

• The Seam Interface: computation_client.h, Line by Line — Everything PyTorch does on an XLA device leaves the process through one C++ class. Walk the real computation_client.h top to bottom: the four different things called "Computation", the transfer surface and its GIL warning, Compile, ExecuteReplicated, and the 43 pure-virtual methods that are the entire seam between PyTorch and PJRT.
• The Lazy Tensor: What Happens Between Your Op and sync() — On the xla device an op runs nothing. It appends a node to a graph, and the tensor you get back is a promise. This walks the whole mechanism in PyTorch/XLA at a pinned commit: what an XLATensor actually holds, how each IR node hashes itself as it is built and why shape only enters through the leaves, what torch_xla.sync() sweeps up and folds into one graph hash, the four ingredients of that hash in the order the executor folds them, the 2048-entry compile cache a hit or a miss lands in, and the reads (.item(), .cpu(), a stray print) that cut the graph somewhere you did not ask. Then the same machinery seen twice more: eager mode as a cut after every op, and torch.compile with the openxla backend as a cut pinned to the function boundary, traced once and replayed by hash.
• How PJRT_DEVICE Becomes a Client: pjrt_registry.cpp, Line by Line — Set PJRT_DEVICE=TPU, import torch_xla, and a chip you never named starts running your model. The whole decision is 174 lines of C++ in one file. We read all of it: the plugin interface with its three questions, the global map seeded with a single placeholder entry, the exact-string lookup whose error message doubles as the documentation, and then the if/else chain itself, branch by branch. The dynamic-plugin branch that a plain import turns on by default, and which dlopens a vendor .so, initializes it, and wraps it through the PJRT C API. The distributed key-value store that only exists for plugins that ask for it, built out of the same coordinator torch_xla uses for preemption. The CPU branch, the one device with no plugin at all. The TPU branch and its three-deep search for libtpu.so. The two half-finished branches, XPU and NEURON, that skip the plugin initialize and the profiler hook. And the else that catches your typo. The signature exhibit runs the real branch order: type a device string, flip the dynamic-plugin switch, and watch which branch catches it and which environment variables it reads.
• The PJRT Boundary: One Training Step, Crossing by Crossing — Somewhere between your PyTorch code and a TPU there is a line, and it is one C++ class: ComputationClient, pure virtual, forty-three methods, one live implementation. This page walks a single lazy training step across it four times (inputs down, program down, run, one value up), names the PJRT method waiting on the other side of each crossing, and counts them against a published lab capture on a Colab TPU. Then the things that never cross: an all-reduce is an instruction inside the program, not a call, and buffer donation rides down as an HLO annotation because the execute options struct is two booleans wide. Closes with the metrics report that has the seam running down its middle, the fourteen files that reach the client, the SPMD path that swaps one method for its plural, and the pinned OpenXLA commit that makes every line number checkable.
• Writing torch_xla Notebooks That Survive Colab — A Colab notebook that trains a model on a TPU is easy to get running once and surprisingly hard to get running twice. This is the craft layer under torch_xla: why torch and torch_xla are one pinned pair and never two independent versions; what PJRT_DEVICE actually does when you leave it unset (import picks a default and tells you in a warning line most people scroll past); why torchax and torch_xla cannot share a runtime, at run time over PyTorch’s dispatcher and again at install time over libtpu; why the first step is slow and the second one is not; and the three instruments that let you prove any of it: the metrics report, the IR and HLO dumps written to a file whose name is not the name you gave it, and the two calls that print a graph without running it. Closes with a cell order that works and a paste-back discipline that keeps a reference run honest.
• torch.compile Meets the Lazy Tensor: dynamo_bridge.py, Line by Line — Name openxla as your torch.compile backend and PyTorch hands the captured FX graph to one 794-line Python file. The obvious guess is that it lowers those nodes to XLA directly. It does not. It runs the graph once, on your real device tensors, through the same lazy runtime an un-compiled program uses, takes the hash of the recording, compiles under that hash without executing, restores every tensor the run modified, deletes the recording, and hands back a closure that holds the hash. We read all of it: the matcher that rebuilds the parameter list from your arguments and the trace-time weights, the three small classes that put duplicated, pass-through and None outputs back before the caller sees them, the tracing function and its five undo steps, the closure that runs on every call and never compiles, the collector that discovers unsupported operations by running the graph node by node and watching a counter, and the partitioner that quietly turns one compiled region into several XLA programs with host code between them. The signature exhibit runs the file’s own branch: press a call, watch the numeric cache key decide whether the trace runs, and watch the launch happen either way.
• Where torch_xla Calls PJRT: pjrt_computation_client.cpp, Line by Line — Everything PyTorch/XLA ever asks a device to do goes through one 1,073-line C++ file, and this walk reads all of it. Initialize, which runs once per process and quietly decides the names your devices will answer to. TransferToDevice, where a host tensor becomes a device buffer and an empty lambda is the only thing keeping the source alive. Compile, whose 143 lines are almost entirely two arms of an if: one that sets num_partitions to the device count and one that sets num_replicas to it, with a device assignment matrix transposed between them. ExecuteComputation and ExecuteReplicated, which take the same executable, disagree about strict shape checking, and lock in completely different ways. TransferFromDevice, the one call that genuinely blocks. And a first function in the file that is never called at all. The signature exhibit runs one step through the file in its own order, each crossing showing the PJRT method underneath it and the timer it stamps.
• SPMD in torch_xla: A Mesh, an Annotation, and One Virtual Device — Your model outgrew one chip. The usual PyTorch answer is to rewrite it into a parallel model; the SPMD answer is to keep writing single-device code and tell the compiler how a few tensors are cut across the chips. This walks the whole path in PyTorch/XLA at a pinned commit: what use_spmd() switches on and why every tensor then lands on a device called SPMD:0, what a Mesh and a partition spec really are, the permutation that turns a spec into a tile assignment, what mark_sharding does to the bytes you already uploaded, how the annotation reaches the HLO, and why one step becomes ExecuteReplicated behind a single device lock.
• From Pending IR to a Device Buffer: xla_graph_executor.cpp, Line by Line — On an XLA device your operations do not run when you write them. They pile up as unexecuted nodes, and then one call turns the pile into a single number. This is the 1,609-line C++ file that does it, read end to end. The compilation cache and the two environment variables that size it. The arena that knows every live tensor, and why a random seed has to be a small graph of its own. The six merge sites, scattered across nine hundred lines, that assemble the graph hash: a config flag a print sets differently from a sync, the two git revisions baked in at build time, one hash per synced tensor, the parameter order, the donated buffers, and the sharding mode. The lookup that decides compile or replay, and the two counters that tell you which happened. Buffer donation and the one condition under which it is safe, argued in the file’s own twenty-line counterexample. Then lowering, sharding annotations, parameter wrapping above 3200 inputs, and the single call that costs the seconds. The signature exhibit folds the real hash term by term, in the order the code folds it, so you can watch where one changed ingredient makes every later digest diverge.