The array moves up

PJRT's unit of data is the per-device buffer: a PjRtBuffer holds one shard's worth of memory on one device, and a single logical array spread across eight devices is eight separate PjRtBuffer objects, with the framework above keeping track of how they fit together. IFRT (xla/python/ifrt/) exists to erase that bookkeeping. ifrt::Array (xla/python/ifrt/array.h:65) is one object: it carries a dtype(), a shape(), and a sharding() of its own, and an eight-device array is one ifrt::Array, not eight of anything.

The array moved up the stack; the buffers stayed behind.

ifrt::Client is what builds these objects. MakeArrayFromHostBuffer turns a plain host buffer into one; AssembleArrayFromSingleDeviceArrays does the adjacent job in reverse, taking per-device shards a caller already has and gluing them into a single logical array. Where PJRT asks the framework to track eight buffers and a sharding by hand, IFRT asks for one array that already knows its own sharding.

Four methods carry the whole idea: dtype(), shape(), sharding(), and layout(), each a pure virtual with no default implementation. An Array cannot be copied or moved either, on purpose: every reference to one is a tsl::RCReference<Array>, aliased as ArrayRef, a ref-counted handle rather than a value that could get duplicated and quietly drift out of sync with its own shards.

The rest of the interface is about taking that one object apart or putting it back together without touching shard data unnecessarily. DisassembleIntoSingleDeviceArrays is the exact inverse of AssembleArrayFromSingleDeviceArrays; FullyReplicatedShard is a shortcut for the common case where every shard is identical, one shard read instead of all of them. An ArrayCopySemantics enum, kAlwaysCopy, kReuseInput, kDonateInput, governs every one of these calls the same way donation governs a PJRT buffer: some operations promise a fresh copy, others let the caller trade ownership of the input away for speed.

class Array : public RTTIExtends<Array, Value> {
 public:
  Array() = default;

  // Not copyable or movable.
  Array(const Array&) = delete;
  Array(Array&&) = delete;
  Array& operator=(const Array&) = delete;
  Array& operator=(Array&&) = delete;

  virtual DType dtype() const = 0;
  virtual const Shape& shape() const = 0;
  virtual const Sharding& sharding() const = 0;

IFRT is an interface, not a runtime, and more than one thing implements it. The PJRT-backed implementation wraps PjRtClient and PjRtBuffer objects directly, in the same process as the Python that called it. The line dividing the two implementations is where the runtime lives relative to your process, not how many hosts the job spans: the multi-controller jobs of the next chapter still use this in-process implementation, one client per host, each wrapping that host's local PJRT client. The proxy implementation, under xla/python/ifrt_proxy/, exists for when client and runtime do not share a process: the process holding the Python ifrt::Client talks over a wire to a server process that owns the actual runtime, wherever that runtime happens to live.

Chapter 13 depends entirely on that split existing. Swapping what sits behind the proxy's server, from an ordinary PJRT-backed runtime to something like Pathways, changes nothing about the JAX code running above the client. The abstraction is not decoration; it is the exact seam a single-controller runtime needs in order to be a drop-in replacement underneath code nobody had to touch.

A jax.Array sits on an ifrt::Array underneath, and that is the concrete reason a single jax.Array can span multiple host processes in JAX's multi-process world, the training run the JAX path meets from above (at /jax/training-run). The sharding lives on the array object itself, not stitched together by the framework from a pile of per-device buffers it has to track by hand.

Client exposes operations that only make sense once an array is a first-class object. RemapArrays shuffles shards between arrays and BitcastArrays reinterprets their bytes; the header calls both metadata-only, with no shard data copied. ReshardArrays changes an array's sharding, and that one can move data, because a shard whose new sharding places it on another device has to be copied there. None of this replaces PJRT. Every IFRT array still bottoms out in real per-device memory, and the PJRT-backed implementation is proof that IFRT is a layer built on top, not a rewrite underneath. What moved is where the bookkeeping lives: up, onto an object that already knows what it is, instead of a set of buffers a caller has to remember to keep in sync.