What is an IFRT array?

It is one logical array that lives on many devices at once, and the implementation over PJRT is thin enough to describe in a sentence: a dtype, a shape, a sharding, an optional layout, and a vector of per-device buffers. The sharding says how the logical shape is cut up and which device holds each piece; the vector holds the actual device memory, one entry per addressable shard. Nothing in the object is a copy of the data. Creating one takes buffers that already exist on their devices and records the correspondence, and destroying one releases the references. The class is the C++ object behind a sharded jax.Array, and the interface it implements, IFRT, is the layer a single controller drives thousands of chips through, sitting above PJRT rather than replacing it.

How does IFRT match buffers to devices?

By position, and only against the devices the current process can address. The validator reads the sharding’s addressable device list, requires the buffer vector to have exactly that many entries, and then walks the two lists together by index. Buffer i must sit on device i. The consequence catches people out: handing over the right set of buffers in the wrong order is rejected, because index 0 is compared against device 0 and nothing searches for a match. The error names both devices. The addressable half matters just as much on a multi-host job, where a sharding can span 512 devices while the process holding it addresses only 8, and the buffer vector is expected to be 8 long rather than 512.

What is the difference between kAlwaysCopy, kReuseInput and kDonateInput?

On paper, three different promises: always allocate fresh buffers so mutating the output cannot touch the input, try to share the input buffers, or take the input buffers and leave the source unusable. In practice it depends which function you reach them through. The helper that disassembly and the fully-replicated shortcut both use has one case per semantics and all three return the same buffer, with a TODO above the first saying that always-copy should clone and that the PJRT API has no efficient same-device clone to do it with. The copy path is the one that honours them: an always-copy to the same device really does call the PJRT copy, and a donation across devices nulls out the source slot so the original array is left holding a hole.

Why does copying an IFRT array to the host fail with "Only single-shard is implemented"?

Because the host read is written for one shard and takes exactly one shortcut past that. It first checks whether the sharding has more than one device. If it does not, it proceeds. If it does, it asks whether the sharding is fully replicated, and if so it grabs a single shard through the replicated-shard shortcut and reads that one instead, which is correct precisely because every shard holds the whole array. Anything else, a genuinely sharded array over several devices, returns that error with the device count spliced in. The way round it is to disassemble first and read the shards individually, or to have the runtime assemble a replicated copy.

What does DisassembleIntoSingleDeviceArrays actually do?

It hands back one single-device array per shard and moves no data. It asks the sharding to split the logical shape into per-shard shapes and per-shard shardings, then for each of those builds a new array wrapping exactly one buffer taken from the original vector. The custom layout is carried across. Two details are worth holding on to. The buffer is fetched through the copy-semantics helper, which returns the same buffer for all three semantics, so the shards alias the original array rather than owning fresh memory. And although the caller can ask for all shards or only addressable ones, the split itself is always requested with addressable-shards semantics, with the all-shards request turned into an up-front error when the sharding is not fully addressable.

IFRT Arrays: pjrt_array.cc, Line by Line

A jax.Array spread over 512 chips is one Python object, and underneath it is a list of ordinary single-device buffers plus a description of how they tile the whole. This 696-line C++ file is the class that holds those two things together, and the walk reads all of it. The validator that runs before any array exists, which matches buffers to devices strictly by position and compares only the devices this process can address. The six construction paths, one of which quietly skips the validator. The three copy semantics that the header describes as different and that one function implements identically, with the TODO admitting it. Disassembly, which hands back one array per shard without moving a byte. The read back to the host, which refuses anything but a single shard unless the array is replicated. And a copy path where the same invariant is re-derived against a different device list. The signature exhibit runs the validator itself: pick a sharding, perturb one buffer, and watch which of the five checks refuses it and with which string.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of IFRT Arrays: pjrt_array.cc, Line by Line — built to be understood, not skimmed.

Questions

What is an IFRT array?: It is one logical array that lives on many devices at once, and the implementation over PJRT is thin enough to describe in a sentence: a dtype, a shape, a sharding, an optional layout, and a vector of per-device buffers. The sharding says how the logical shape is cut up and which device holds each piece; the vector holds the actual device memory, one entry per addressable shard. Nothing in the object is a copy of the data. Creating one takes buffers that already exist on their devices and records the correspondence, and destroying one releases the references. The class is the C++ object behind a sharded jax.Array, and the interface it implements, IFRT, is the layer a single controller drives thousands of chips through, sitting above PJRT rather than replacing it.
How does IFRT match buffers to devices?: By position, and only against the devices the current process can address. The validator reads the sharding’s addressable device list, requires the buffer vector to have exactly that many entries, and then walks the two lists together by index. Buffer i must sit on device i. The consequence catches people out: handing over the right set of buffers in the wrong order is rejected, because index 0 is compared against device 0 and nothing searches for a match. The error names both devices. The addressable half matters just as much on a multi-host job, where a sharding can span 512 devices while the process holding it addresses only 8, and the buffer vector is expected to be 8 long rather than 512.
What is the difference between kAlwaysCopy, kReuseInput and kDonateInput?: On paper, three different promises: always allocate fresh buffers so mutating the output cannot touch the input, try to share the input buffers, or take the input buffers and leave the source unusable. In practice it depends which function you reach them through. The helper that disassembly and the fully-replicated shortcut both use has one case per semantics and all three return the same buffer, with a TODO above the first saying that always-copy should clone and that the PJRT API has no efficient same-device clone to do it with. The copy path is the one that honours them: an always-copy to the same device really does call the PJRT copy, and a donation across devices nulls out the source slot so the original array is left holding a hole.
Why does copying an IFRT array to the host fail with "Only single-shard is implemented"?: Because the host read is written for one shard and takes exactly one shortcut past that. It first checks whether the sharding has more than one device. If it does not, it proceeds. If it does, it asks whether the sharding is fully replicated, and if so it grabs a single shard through the replicated-shard shortcut and reads that one instead, which is correct precisely because every shard holds the whole array. Anything else, a genuinely sharded array over several devices, returns that error with the device count spliced in. The way round it is to disassemble first and read the shards individually, or to have the runtime assemble a replicated copy.
What does DisassembleIntoSingleDeviceArrays actually do?: It hands back one single-device array per shard and moves no data. It asks the sharding to split the logical shape into per-shard shapes and per-shard shardings, then for each of those builds a new array wrapping exactly one buffer taken from the original vector. The custom layout is carried across. Two details are worth holding on to. The buffer is fetched through the copy-semantics helper, which returns the same buffer for all three semantics, so the shards alias the original array rather than owning fresh memory. And although the caller can ask for all shards or only addressable ones, the split itself is always requested with addressable-shards semantics, with the all-shards request turned into an up-front error when the sharding is not fully addressable.

IFRT Arrays: pjrt_array.cc, Line by Line

Questions

Related explainers