Bridges

Everything in this path so far runs on whatever device happens to sit under Python: a laptop CPU, the Intel chip this course was written on, maybe a GPU somewhere. None of it runs on a TPU without crossing first, and PyTorch itself does not build that crossing. Two separate projects do, and this chapter is about both of them at once, because knowing only one leaves you unable to read half of what you will find once you go looking.

Say this plainly, once, because it matters more here than it has anywhere earlier in the path: neither this chapter nor the next ran on a TPU. The machine that wrote this course has no TPU attached to it. Every fact below is checked against the current pytorch/xla repository and its official docs rather than printed from a real run, and the Colab lab pass later in this path is where these same facts turn into your own numbers, measured on real hardware. That pass has since been run: the numbers it produced live in chapter 12, and the three failures it produced on the way are exhibits in the museum's pytorch wing.

That gap does not make the facts soft, but it does change what fluency looks like for these two chapters. Instead of predicting a printed value, the goal is reading two bridges well enough to know which one a given module would cross, and what is waiting for it on the other side.

The first bridge is torch_xla, and its central idea is laziness. Move a model and its tensors onto the xla device and every operation you run against them stops executing right away; it gets recorded instead, the way the tape in chapter 2 records an op, except this tape waits for a sync point before anything reaches the chip. torch_xla.step() marks that boundary inside a training loop, and torch_xla.sync() marks it directly: only there does the recorded graph actually materialize, compiled through XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → and run through PJRT, with PJRT_DEVICE=TPU telling the runtime which hardware to reach for.

The classic cost of this bridge is recompilation. Change the shape feeding a lazy graph and the sync point has to compile a new graph for it, the same cache-key discipline chapter 6 taught for dynamo, paid again at a different layer. A training loop whose batch shapes drift, a ragged last batch above all, pays this bill on every drift instead of once. The xla path's chapters 4 through 6 on this site describe exactly what runs underneath that sync point, StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → in and XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →'s own decisions out; this chapter only needs you fluent in when the sync happens and why it costs what it costs.

The second bridge, torchax, began inside the pytorch/xla repository and now lives in its own (google/torchax), and it takes an entirely different route. Instead of recording torch ops for a later sync, it maps them onto jax primitives directly: a torch tensor under torchax is backed by a real jax array underneath, not a lazy recording of one. Because that array genuinely is a jax array, everything the jax path spent nine chapters building, jit, grad, sharding over a named mesh, applies to a torch module exactly as it would to a jax function, with no separate torch-side reimplementation waiting to be written.

The interop runs in both directions, and that detail is worth holding onto. A jax program can call a torch module the way it calls any other jax function, and torch code can reach into jax the same way back. Neither side has to pretend the other framework does not exist; torchax's whole job is making that pretending unnecessary.

Follow either bridge far enough and they stop looking like two things. torch_xla's sync point and torchax's jax-backed tensors both bottom out at StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo →, the same portable tensor IR the xla path's chapter 2 teaches from the jax side, and the kernel path's own account of the stack, at /l/source, describes from below. By the time the compiler is looking at the program, it has no memory of which framework wrote it. The framework identity that felt so load-bearing through chapters 1 through 9, mutation against threading, guard-and-recapture against trace-once, is gone by this layer, and that convergence is exactly why two different bridges can lead to the same chip.

Two bridges, one destination: StableHLO.

torch.export, chapter 7's whole-graph artifact, is exactly the shape both bridges want hold of: no Python left in it, decomposed toward a core set of ATen ops, nothing left implicit for a bridge to guess at. Export a module before you hand it to either bridge and you have already done the part of the crossing that depends on you; what happens after that is the bridge's job, not yours.

Until a real chip is measuring these facts back at you, the work this chapter asks for is legibility, not a benchmark. Given a module, you should be able to say which bridge would carry it and why, what its execution model looks like once it crosses, and what the compiler is left holding once both bridges converge. The torchax README is short enough to read end to end in one sitting, and it is worth doing exactly that before chapter 11, where this same convergence stops being theory and starts being the ground you stand on.