Donating the state tree

What donate_argnums promises for a single argument, what becomes of the array you gave away, and the two mismatches that get a donation refused all belong to the performance chapter's lesson on donation. Read that one first; this lesson starts where the donated argument stops being an array and starts being a tree. What the compiler then does with a granted alias, and where it inserts copies to keep a live value from being clobbered, is 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 → path's at /xla/layout-memory/copies-and-buffers.

Mark the state argument and the header of the compiled module comes back with one entry per leaf: input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias) }. Output position zero may take input position zero's buffer, and so on down the flattened tree. That is the one-leaf-one-slot correspondence from lesson one, read back from the compiler's side.

alias_size_in_bytes totals the map, and the total is checkable against the state itself. It reads 84 here, and sum(leaf.nbytes for leaf in jax.tree.leaves(state)) also reads 84, so every leaf of this state is covered and none was left out. Check that equality on a real parameter tree rather than assuming it, because a state can come back with most of its leaves aliased and one of them not.

from dataclasses import dataclass

import jax
import jax.numpy as jnp

@jax.tree_util.register_dataclass
@dataclass
class State:
    w: jax.Array
    b: jax.Array
    step: jax.Array

def sgd(state, g):
    return State(w=state.w - 0.1 * g, b=state.b + 1.0, step=state.step + 1)

state = State(w=jnp.ones((4, 4)), b=jnp.zeros(4), step=jnp.zeros((), jnp.int32))
c = jax.jit(sgd, donate_argnums=0).lower(state, jnp.ones((4, 4))).compile()
head = c.as_text().splitlines()[0]
print(head[: head.index("entry_computation_layout")])
print(c.memory_analysis().alias_size_in_bytes)
print(sum(leaf.nbytes for leaf in jax.tree.leaves(state)))

# HloModule jit_sgd, is_scheduled=true, input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias) },
# 84
# 84

A donated array either aliases or it does not. A donated tree has a middle case: most leaves keep their alias, one drops out of the map, and the executable ends up doing part of what you asked for.

The two functions below differ in a single character. One returns the step counter as state.step + 1, the other as state.step + 1.0, so an int32[] goes in and a float32[] comes back. The map loses its {2} entry, the total falls from 84 bytes to 80, and the missing 4 is the counter. The weight and bias entries are identical in both.

The refusal also arrives on stderr, in the same warning the performance lesson reads, naming the aval it could not place. What the map adds is arithmetic the warning cannot give you: which leaf, out of how many, and what fraction of the promise survived. A whole-tree answer would say the donation failed. This one says you got 80 of the 84 bytes you asked for.

On a state, donation is not granted or refused. It is granted leaf by leaf.

from dataclasses import dataclass

import jax
import jax.numpy as jnp

@jax.tree_util.register_dataclass
@dataclass
class State:
    w: jax.Array
    b: jax.Array
    step: jax.Array

def kept(state, g):
    return State(w=state.w - 0.1 * g, b=state.b + 1.0, step=state.step + 1)

def drifted(state, g):
    return State(w=state.w - 0.1 * g, b=state.b + 1.0, step=state.step + 1.0)

state = State(w=jnp.ones((4, 4)), b=jnp.zeros(4), step=jnp.zeros((), jnp.int32))
for fn in (kept, drifted):
    c = jax.jit(fn, donate_argnums=0).lower(state, jnp.ones((4, 4))).compile()
    head = c.as_text().splitlines()[0]
    print(fn.__name__, head[head.index("input_output_alias") : head.index("entry_computation")])
    print("   ", c.memory_analysis().alias_size_in_bytes, "bytes")

# kept input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias) },
#     84 bytes
# drifted input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias) },
#     80 bytes

Donating a state donates its leaves, and a leaf is an object, not a location. Pull one out and stash it somewhere else, last epoch's weights in a metrics dict, a slice held back for a diff, and the thing you stashed is the same array the compiler was given permission to overwrite. It goes with the tree.

The run below keeps state.w under a second name, calls one donating step, and asks both names what survived. The kept name reports its array deleted; the batch, which was never donated, reports False. So the audit before you add the keyword is not about the step at all. It is about every other place in the loop that still holds a leaf.

That also settles which argument takes the keyword. The state is replaced by the step's own output, so nothing upstream needs it once the call returns. The batch arrives from the host and downstream code usually still wants it for a metric, which makes donating it a way to delete an array you are about to read.

from dataclasses import dataclass

import jax
import jax.numpy as jnp

@jax.tree_util.register_dataclass
@dataclass
class State:
    w: jax.Array
    b: jax.Array

def step(state, batch):
    return State(w=state.w - 0.1 * batch, b=state.b + 1.0)

fast = jax.jit(step, donate_argnums=0)

state = State(w=jnp.ones((4, 4)), b=jnp.zeros(4))
batch = jnp.ones((4, 4))
kept = state.w                    # one leaf, held under a second name
state = fast(state, batch)
print("the leaf held elsewhere:", kept.is_deleted())
print("the batch:", batch.is_deleted())
print("the new weight:", state.w[0, 0])

# the leaf held elsewhere: True
# the batch: False
# the new weight: 0.9