[blasphemetheus]

Adds checkpoint/2 to Nx.Defn.Kernel. Closes #765.

defn block(x, w1, w2) do
  checkpoint([x, w1, w2], fn x, w1, w2 ->
    x |> Nx.dot(w1) |> Nx.relu() |> Nx.dot(w2)
  end)
end

Reverse-mode differentiation keeps every intermediate of the forward pass until the backward pass has used it. Inside a checkpoint the backward pass instead reconstructs the body's intermediates by running the body again from its inputs, so only the inputs and outputs of the block stay alive across the forward pass. Outside grad a checkpoint is the same as calling the function.

How it works

checkpoint/2 builds an Nx.block carrying a %Nx.Block.Checkpoint{} struct, with the inputs passed as one container argument. The gradient rule for that block:

1. passes the saved inputs and the incoming gradients through one barrier (a barrier: true metadata node, lowered to an optimization barrier by EXLA),
2. calls the body again on the barrier's outputs,
3. differentiates the fresh copy with respect to the block inputs, including any cond inside the body, and feeds the input gradients back into the outer traversal.

The barrier is what makes the memory saving possible. Without it XLA deduplicates the recomputed body against the forward body, and with the inputs alone behind it XLA may schedule the recomputation right after the forward pass. Tying the incoming gradients into the same barrier delays the recomputation until the backward pass reaches the block. JAX's checkpoint works the same way.

The gradient's value does not depend on the compiler. The memory saving does: it needs the compiler to keep the recomputation separate and schedule it late. EXLA does, on GPU and, since XLA 0.11, on CPU.

Limits

• inputs is a tensor or container, or a list of them for a function of several arguments; the body must return a tensor or a flat tuple of tensors.
• The body cannot read tensors from the enclosing scope. Capturing them automatically is a possible follow-up once the design is agreed.
• The body may run more than once, so random numbers inside it need their Nx.Random key passed as an input. There is a test for that.
• Recomputation is all or nothing per block. Selective policies would be an option on the same function and a field on the same struct.

Tests

nx/test/nx/defn/checkpoint_test.exs covers correctness against the plain function: containers, vectorization, cond, while, custom_grad, shared weights, nested and back-to-back checkpoints, higher-order gradients through nested checkpoints with linear and nonlinear conditionals, duplicated outputs, a random mask from an explicit key, and recomputation counts. exla/test/exla/defn/checkpoint_test.exs checks the compiled gradient keeps the recomputed body, recomputes only what it needs, and lowers peak scratch memory on an MLP.

🤖 Generated with Claude Code

https://claude.ai/code/session_01PpGtcRtwc1jfXLCeqVzs6s

[blasphemetheus]

Figured this gives fuller context to some of the decisions in #1862 , the envisioned api would only be possible via macros I believe which will be in a draft PR #1866 can be done with lifting i think

[blasphemetheus]

If there is a viable path that diverges sharply from how this is implemented, I would close the PRs that diverge and stack PRs from the point that is used

[blasphemetheus]

to make this one a PR stack it'd need the branch that the PR comes from to be from this repo's branches, not a forks I believe. So to get Github Stacks to work there's a server-side object linking PRs, just doesn't cross the fork boundary

[blasphemetheus]

found a couple instances of weird behavior regarding differentiation, which checkpointing can run into, but isn't caused by this PR stack in either case, I'll make reproducer tests for the weird behavior and leave the original failing tests in this one, can merge after fixes if the fixes are straightforward

[blasphemetheus]

hmm I force-pushed a bunch of new commits with new names, which is probably less helpful than just pushing new commits on top of all this.

• b0fe1e7c is the new stuff, touching grad.ex in the checkpoint gradient rule

• added a test in 3ad0d54b

[blasphemetheus]

So this PR uses block to pass tensors, like: qr, Take, and TopK do. But. The body is written by the user, so it can read variables from the surrounding defn that weren't passed into the block.

Right now when that happens, there is no indication of it, The forward pass returns an unevaluated expression on the Evaluator, EXLA fails to compile the block, and the gradient of the outside tensor comes out as zeros. The current failing tests should show this behavior.

One way to fix this could be to give a block body it's own context, like the way while has its own context, so mixing in an outside tensor raises at trace time.

[blasphemetheus]

#1895 seems like it will make this PR less complicated if it is merged, so I'll simplify this one if that direction goes forward

[blasphemetheus]

new stuff is container arguments for Nx.block replacing the arity limit (capped at 253 before)

I also trimmed some of the tests that seemed overlong, specifically the ones that only varied the body, kept all the more structure-based variations

Having outputs be accepted as maps could be a future direction. The issue also mentions selective recomputation policies as a future direction. Neither are tackled in this PR