Above the compiler

IFRT's own README is unusually direct about why it exists, and it is worth reading before any diagram. IFRT stands for Interim Framework Runtime, and it is described as a high-level ML runtime API designed to be used as the interface between a user-facing framework, such as JAX, PyTorch, or TensorFlow, and the runtimes below.

The history is stated just as plainly. PjRt was a low-level API that abstracted hardware differences on a single host, and it was, in the README's own phrasing, pretty much just another runtime. Distributed execution across thousands of accelerators went past what that API was shaped for, so the teams decided to bifurcate the current PjRt API and let the two deviate to better support their intended use cases. IFRT is not a rewrite of PJRT and it does not replace it.

One sentence in there is the design brief for everything else in this lesson: frameworks should more or less declaratively express the work that needs to be done, and delegate policy choices about how to efficiently execute that work to the runtime implementations. Declare what, not how. Every difference between the two APIs follows from taking that seriously.

The chapter above makes the central comparison and quotes array.h directly: PJRT's unit is a per-device buffer, so an eight-device array is eight PjRtBuffer objects with the framework tracking how they fit together, while ifrt::Array is one object carrying its own dtype, shape, sharding, and layout. That is the correct headline. What a chapter cannot do is show you the rest of the directory, and the rest of the directory is where the second half of the story is.

The headers in xla/python/ifrt include array.h, client.h, device.h, device_list.h, executable.h, compiler.h, program.h, dtype.h, layout.h, memory.h, and bundle.h. Two of those names say something the array comparison alone does not. compiler.h and program.h mean IFRT has its own notion of compilation, so a runtime behind it can be handed a program to compile rather than only arrays to move. device_list.h means a set of devices is a first-class value with its own type, which PJRT has no equivalent of.

Put those together and the layer is doing more than bookkeeping. A framework can hand IFRT a program and a set of devices and a sharding, and let the implementation decide how to place and execute it. That is the declarative brief from the README, expressed as a header list.

array.h        one logical array: dtype(), shape(), sharding(), layout()
client.h       what builds and reshapes those arrays
device.h       a device
device_list.h  a SET of devices, as a value type; PJRT has no counterpart
compiler.h     IFRT compiles, it does not only move data
program.h      ... and it has its own notion of what a program is
executable.h   the compiled result, across devices
dtype.h  layout.h  memory.h  bundle.h

An interface only earns the name if something other than the obvious implementation can sit behind it, and IFRT has two. The in-process implementation wraps PjRtClient and PjRtBuffer directly, and a single-process JAX program runs on it without ever mentioning IFRT. The other one, under xla/python/ifrt_proxy, splits the interface in half: the process holding the Python client talks over a wire to a server that owns the actual runtime.

That split is the entire swap mechanism, and /xla/pathways is where the path spends a chapter on what gets swapped in. Pathways puts a single controller in front of a fleet of workers, which buys MPMD, worker failure as a recoverable event rather than a job-ending one, and a scale ceiling above a single gang-scheduled program. Its design is published as a paper. Its runtime is not open source, and the honest version of this lesson says so in the same breath: what you can read is the proxy boundary, because that boundary has to be public for anything to plug into it at all.

So the property that makes the layering worth its complexity is testable in one sentence. The same JAX program, unchanged, runs against an in-process PJRT-backed IFRT client and against a proxy client with a different runtime behind it. Nothing above the client had to know which.

It is fair to ask whether this is one abstraction too many, and the answer is easiest to see by trying to collapse it in each direction. Fold IFRT into PJRT and every framework goes back to tracking eight buffers and a sharding by hand, which is exactly the bookkeeping JAX had before and exactly the thing that does not survive a single logical array spanning host processes.

Fold the runtime into IFRT and the interface stops being an interface. The proxy split only means something because there is a real seam there with two implementations on the far side of it, and a single-controller runtime is not a variant of an in-process client, it is a different system with a different failure model.

Which leaves three layers with one job each. PJRT owns per-device memory and one compiled executable on one client. IFRT owns arrays, shardings, device sets, and a compilation interface across devices. The runtime behind the proxy owns scheduling and placement policy. Each of them is the smallest thing that can answer its own question, and none of them can answer another's.

These lessons went down the stack and back up it: the exported symbol at pjrt, the invariants a pass may not break at hlo, the pipeline as filenames, the fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → queue, the partitioner, the seam where backends part ways, and the layer above all of it here. What is left is not more depth pages. It is the path's own closing chapters and the work in them.

/xla/ifrt reads array.h line by line, /xla/mcjax has the multi-controller architecture these arrays span, and /xla/pathways states carefully which claims come from a paper rather than from code. Read them in that order if you have not, or reread them now that the interfaces underneath have names.

Then /xla/capstone, and specifically its first project, which is the one these lessons were assembled to make possible: implement the minimal PJRT C API surface until jax.devices() on your own machine returns a device you built. The pjrt lesson already listed what CreatePjrtApi will ask you for. The deliverable that project wants is a write-up of every struct you filled in and the order you filled them, which is a map nobody has published and which you will only have after doing it.