Two kinds of copies, and who owns a buffer

The layouts lesson under the XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → layer chapter owns the reading of {1,0} annotations and the copies a disagreement inserts; this lesson picks up where that one stops, at position and ownership. Layout assignment has an address in the pipeline: in this repo's CPU capture it sits mid-flight, after flatten-call-graph and before sub-byte-size-setter, which means everything after it must either respect the chosen layouts or pay for a rearrangement. A layout copy is that payment, and it is attached to a disagreement you can name: two ops, one value, two opinions about which dimension varies fastest.

Late in the same capture, one step before dead-code elimination, runs a pass called copy-insertion, and it has nothing to do with layouts. Its job is protecting live values under aliasing. A while loop's state buffer gets updated in place across iterations; an output that reuses an input's buffer overwrites it; and in every such case, if someone still needs the old value after the new one lands, the program is wrong. The pass finds those hazards and inserts a copy exactly where an alias would corrupt a value that still has a reader.

So the two copy families answer different questions. A layout copy answers "these two ops disagree about memory order." A correctness copy answers "this buffer is about to be reused while its old contents still matter." When a dump shows a copy you did not expect, deciding which family it belongs to is the first diagnostic step, because the fixes live in different places: layouts are argued with the compiler, aliasing hazards are argued with your own program structure.

Aliasing is not only something the compiler discovers; you can hand it over deliberately. jax.jit(f, donate_argnums=...) tells XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → an input's buffer may be consumed: the output can be written into the donated input's memory, saving an allocation and often a copy, which matters exactly when the arrays are large and the update is the point, an optimizer state, a KV cache, a params tree. The price is that the donated array is dead to your program afterward; touch it again and you get an error rather than stale data.

Donation is a request, not a command. A backend that cannot honor it ignores it, and jax warns when that happens, so the honest workflow is to donate, then check the warning stream and the memory numbers rather than assuming the reuse happened. Between this and the correctness copies above, buffer ownership stops being folklore: you can say who owns a region of memory at any point in the compiled program, and what evidence backs the claim.