TPU practice

The disclosure from chapter 10 holds here too, so say it once and move on: nothing in this chapter ran on a TPU either, for the same reason, and the same Colab lab pass is where it turns from a checked fact into a measured one.

What does change once you actually cross is which machine is underneath you, and that machine does not know which framework it is running. Mark shardings on a torch tensor over a named mesh through torchax, the same NamedSharding and PartitionSpec objects the jax path's chapter 10 built, and GSPMD partitions the torch program exactly the way it partitions a jax one: same collectives inserted, same propagation, same cost model. Nothing about SPMD had to be reinvented for torch; it only had to be reached through a bridge that already speaks jax underneath.

The TPU does not care which framework you came from.

Profile a torch model running on a TPU and you land in the same XProf world the kernel path's gym reads live at /gym/kernels#timeline: a capture, a device plane, hardware time accounted for op by op. The profiler was built once, for one machine, and it has no separate torch-shaped version of itself. Whatever discipline chapter 9 taught for reading a trace on CPU or GPU carries over here unchanged; only the device plane you are looking at is TPU instead of CPU.

Pallas kernels, the whole subject of the kernel path on this site, are callable directly from torch code running on TPU. That fact is worth sitting with, because it means a custom kernel was never really a PyTorch question or a JAX question. It is a question about which side of the compiler boundary you are willing to write by hand, and once you are willing, the framework you started in stops mattering.

The clearest proof that this bridge story is real architecture and not a demo already sits in production. vLLM's TPU backend, tpu-inference, runs PyTorch model definitions through a jax lowering path, the same convergence chapter 10 described, serving real models at real scale without asking the model's author to rewrite anything. If the bridge were a toy, nobody would be routing inference traffic through it.

The work this chapter asks for is a map, not a run. Take tpu-inference's stack apart in writing and place every piece of it on the chapter of this site that actually teaches the layer underneath: the SPMD machinery on the jax path's chapter 10, the kernels on the kernel path, the profiler on this chapter. Shard one of your own matmuls over a mesh on paper first, name every collective GSPMD would have to insert, and check your answer against the jax path's chapter 7 before you trust it.