Single controller

Pathways inverts the shape chapter 12 built. Instead of one Python process per host, each blindly trusting the others to be in lockstep, there is exactly one Python client process, full stop. Underneath it sit a resource manager and a fleet of per-host workers that execute whatever programs the client hands them, gang-scheduled across however many pods the job spans.

The design is not a product improvement bolted onto McJAX; it is a different paper's answer to the same problem. The Pathways paper describes an asynchronous distributed dataflow system with centralized scheduling, amortized over large operations so that one controller issuing instructions to thousands of chips does not become the bottleneck it sounds like it should be.

The place chapter 11 spent an entire chapter building an abstraction is exactly the place Pathways plugs in. The client speaks ifrt_proxy, a client and server split of the IFRT interface, to a server that fronts the Pathways runtime on the other side. Your JAX code above that boundary does not change at all: the same jax.jit, the same arrays, the same sharding specs from chapter seven, now dispatching across a proxy instead of an in-process client.

One controller, thousands of chips, the same JAX.

That is the direct answer to a question worth asking plainly: how do you swap out the entire runtime underneath a model without touching the model. You do not touch the model. You touch the client IFRT talks to, and IFRT was built with exactly that seam in mind.

Set JAX_PLATFORMS=proxy and JAX's backend becomes the open-source IFRT proxy client, the wire half of chapter 11's split. The product documentation names what sits on the other side, component by component: a proxy server, a gRPC front that receives the client's requests, and behind it a component the docs call the Pathways client, described there as an IFRT implementation in its own right, one that receives HLO programs and works with a resource manager to place them. Read that list twice and the shape appears: two IFRT implementations stacked in one chain, the open proxy pair doing transport, the closed Pathways one behind it doing the work.

The resource manager is the piece McJAX never had. It runs on plain CPUs, owns allocation across every worker, monitors their health, pauses and resumes jobs, and serves as the single place errors surface. Chapter 12's architecture had no process that could play this role, because every process was busy being a peer.

The paper fills in what the product page abstracts away. A traced program becomes a location-agnostic intermediate representation, a custom MLIR dialect, in which each compiled function is one node in a dataflow graph; a tracer can wrap a Python block that calls several jitted functions and capture the whole block as one program, which is what makes MPMD expressible at all. Buffers enter that IR through a sharded-buffer abstraction, one logical array distributed over many devices, bookkept once. Chapter 11 built exactly this instinct into IFRT, and the resemblance runs in that direction, paper first.

Placement is a negotiation with the resource manager. The client asks for virtual devices, optionally constrained by type, location, or interconnect topology; the manager maps virtual onto physical, and the IR is lowered until it carries real device locations plus explicit transfer operations, scatters and gathers, between computation shards. The lowered program then becomes a dataflow program on PLAQUE, a closed-source production sharded dataflow system that carries all cross-host coordination over the data center network.

Execution is gang-scheduled per island: one centralized scheduler consistently orders every computation on its island, so two programs contending for the same chips cannot deadlock holding half each. Each host runs an executor and a sharded object store, an HBMThe chip’s main memory: large, far, and the resource memory-bound ops spend. 8.2e11 bytes per second on v5e, 1.6e12 on v6e.taught in /l/tpu →-aware cousin of Ray's, holding buffers behind opaque handles the system can migrate. And the controller does not wait its turn per operation. A compiled function's resource needs are statically known, so host-side setup for a successor runs before its predecessor finishes, and one message describing an entire subgraph lets the scheduler sequence all of its shards back to back. That is the paper's answer to the objection its own design invites, one controller in front of thousands of chips, and it is the reason the system is asynchronous dataflow rather than remote procedure calls.

Centralizing control is not free, but what it buys is real. MPMD becomes possible: different programs running on different islands of the same job, the shape pipelining needs. The controller also outlives any individual worker, so a worker failing is a recoverable event instead of a gang-wide one, the elasticity chapter 12 explicitly could not offer. And the scale ceiling moves: a single gang-scheduled program is no longer the largest unit of work the system can express, because one controller can orchestrate many gangs across many pods at once.

The product documentation gives the worker one sentence of contract: a process on a TPU VM that receives compiled executables and performs the computations. Sit with the first half of that sentence, because it inverts chapter 12. Executables arrive at the worker already compiled; compilation happens once, in the head components, instead of once per host across the whole job, and McJAX's thousand-identical-compiles problem simply does not exist here. A sidecar gRPC server on the same VM rounds out the picture, running user-supplied Python next to the chips so data-adjacent work skips the round trip through the controller.

What launches those executables is the part no public source names. Whatever it is must do a PJRT client's exact job, load a compiled program, hold device buffers, fire execution through the TPU runtime, and the paper says Pathways builds on XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → to represent and execute TPU computations. Whether that layer speaks the literal PJRT C API or Google-internal machinery that predates it is stated nowhere you can read, and this track will not guess. What is certain is that the XLA compiler never left the loop: IFRT receives HLO programs, workers receive compiled executables, and something in between is running the same passes chapter 4 dumped. Pathways replaced the runtime and the coordination, not the codegen.

Here is the part worth stating without hedging: the Pathways runtime itself is not open source. You can read the paper that describes its design, you can read the ifrt_proxy client and server code in the public XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → tree, since that boundary has to be public for anything to plug into it, and you can read the product documentation for the cloud surface built on top of it. Be precise about what that public server code is, though: scaffolding, a wire front that serves whichever ifrt::Client it is handed. The proxy server a Pathways cluster runs stands at the same boundary and speaks the same protocol to your client, but whether that deployed binary shares a line of code with the public one is not knowable from the tree, and the runtime it fronts is closed either way, the same way the kernel path marks libtpu closed rather than pretending otherwise. Treat every claim in this chapter about the runtime's internals as coming from the paper and the product surface, not from code you can read yourself.