[blasphemetheus]

Turns checkpoint/2 into a macro that passes tensors read from the enclosing scope into the checkpoint, and adds checkpoint/1 for a body that takes everything from scope.

defn layer(x, w, b) do
  checkpoint(x, fn x -> Nx.dot(x, w) + b end)
end

Before this change every tensor the body used had to appear in inputs, so each layer of a model listed its weights twice: once as an input and once as a parameter of the function. Now only the activation is listed. The gradient of a captured tensor flows through the checkpoint like that of any other input.

How it works

The macro looks at the inline fn, collects the variables its body reads that are bound in the caller and not by the function head, and appends them as extra parameters of the function and extra inputs of the block. A variable the body rebinds under a name from the caller is collected too; that only adds an unused input. A body that is not written inline, such as &layer/3, cannot be inspected, so nothing is captured and every tensor must be passed in inputs. checkpoint/1 raises at compile time for that case.

Nx.Block.Checkpoint gains a saved field listing the positions of the explicit inputs. The gradient rule ties only those to the incoming gradient through the barrier. Captured tensors are read directly, because they are usually parameters that are alive for the whole program, so tying them buys nothing. With no explicit inputs every tensor is saved.

Why captures are read directly

The barrier in the gradient rule exists to keep the recomputation from being scheduled early or merged with the forward pass. For that it is enough that the activations and the incoming gradient go through it; the recomputation cannot start until both are available. Pushing the weights through the same barrier would add nothing, and keeping them out leaves the compiler free to place their reads wherever it wants.

Tests

The existing tests pass unchanged, since a function with no captures compiles to the same block as before. New tests in nx/test/nx/defn/checkpoint_test.exs under "captured tensors":

• captured weights, a captured map of parameters, and checkpoint/1 compute the same gradient as the plain function
• a captured tensor that is the grad target, and one that is not
• the body may rebind a name from the enclosing scope
• indexing a captured variable passes it whole, checked on the expression tree
• checkpoint/1 with a non-inline function raises at compile time
• captured vectorized tensors keep their axes in the forward pass and in the gradient, including a vectorized constant built inside the function

Two tests in exla/test/exla/defn/checkpoint_test.exs check a captured weight and a captured vectorized tensor under EXLA.

Stacked on #1865, which stacks on #1860 and #1862. The diff shows all of them until they merge; the change itself is the macro in Nx.Defn.Kernel, the saved field and its use in the gradient rule, and the new tests, as reviewed on the fork draft blasphemetheus#25. Full nx suite and the EXLA checkpoint tests on CUDA pass.

🤖 Generated with Claude Code

https://claude.ai/code/session_01PpGtcRtwc1jfXLCeqVzs6s

[blasphemetheus]

wew, ok that's the last of the stacked PRs for the version that I've made so far. This one is to get the passing a function to work correctly and adds tests for that api.

[josevalim]

I don't think this is a good design for a macro because it makes me think this is doing to work:

callback = fn x -> Nx.dot(x, w) + b end
checkpoint(x, callback)

But it won't. Ideally we don't rely on macros here. Perhaps we can reuse the block API and either add an option or use Expr.metadata to annotate the block with checkpoint information.

[blasphemetheus]

closing this one as diverging and not the right way