PJRT

JAX, TensorFlow, and PyTorch/XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → are three different frontends with three different programming models, and none of them talk to XLA's optimizer directly. Each one calls into PJRT instead, a runtime API that sits between whatever language or framework produced a program and the compiler that turns it into a running executable. Learn this boundary first and the rest of the compiler stops feeling like one undifferentiated block: everything below PJRT is XLA's business, and everything above it is the frontend's.

That is also why this chapter comes first here. The JAX track ended its own path right above this line, at the point where a lowered function became a StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → string ready to hand off. This track picks up exactly where that one stopped, because PJRT is the boundary every later chapter keeps returning to: it is where a device gets named, where a compiled program gets loaded, and where a result actually comes back.

Every framework speaks to XLA through the same door.

PjRtClient is the object that owns everything else: it holds the list of devices, knows the topology connecting them, and is what you call Compile on to turn a StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → module into something runnable. One process typically holds one client per backend, and every other PJRT object you touch traces back to it.

PjRtBuffer is a handle to memory resident on one specific device, nothing more. It gets created either by handing the client a host array or by running an executable and catching its output. Crucially, a buffer lives on exactly one device: there is no such thing as a buffer that spans two chips at this layer.

PjRtLoadedExecutable is a compiled program bound to a specific set of devices and ready to run. Execute takes a nested list of buffers, one inner list per partition, and returns the same shape back. Client owns devices and compilation, buffers are per-device memory, executables run: three nouns, three jobs, and almost everything else in PJRT is plumbing between them.

There are two ways a frontend actually calls into PjRtClient. The first is in-process C++: PjRtStreamExecutorClient links directly into the same binary as the frontend and gets called through ordinary virtual methods. This is the simpler path, and it is how XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → talked to its callers for years.

The second is the C ABI, a plain struct of function pointers defined once and versioned carefully. A backend that wants to ship without merging a single line into 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 source builds a shared library exporting that struct, and the frontend loads it at runtime rather than linking it at build time. This is not a hypothetical: Intel's oneAPI backend, Apple's Metal backend, and AWS Trainium all ship as exactly this kind of plugin.

It is easy to assume a new accelerator has to get its support baked into 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 repository before anything works. The C ABI exists specifically so that is false: the struct is the entire contract, and a vendor who honors it never needs XLA's build system, its review queue, or even its source code.

Compile takes a StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → module in. Execute takes buffers in and returns buffers out. That is the entire data-plane contract PJRT enforces, and it is worth sitting with what it implies about sharded arrays: one logical array spread across eight devices is not one PjRtBuffer with a fancy shape. It is eight separate PjRtBuffers, one per device, and PJRT itself has no notion that they belong to the same array at all. Chapter 11 is where a layer above PJRT gives that array a single identity again.

On a real machine the objects are easy to see without reading a line of C++. Ask a device for its platform, its kind, and the type of the client underneath it, and PJRT is right there, one attribute access away.

import jax

dev = jax.devices()[0]
print(dev.platform, dev.device_kind)   # cpu cpu  (this machine)
print(type(dev.client).__name__)       # Client: the PJRT client underneath