[blasphemetheus]

Nx.block/4 now accepts a container of tensors, such as a tuple or a map, as an argument.

A block's default implementation receives the struct and then one argument per input, and an Erlang function cannot take more than 255 arguments. A block with hundreds of inputs, such as a checkpointed group of layers whose weights all arrive as block arguments, could not be written. With containers, those inputs are passed as a single tuple or map, and the callback keeps its one-argument-per-input shape:

Nx.block(%MyBlock{}, [x, %{scale: scale, bias: bias}], nil, fn _, x, %{scale: scale, bias: bias} ->
  Nx.add(Nx.multiply(x, scale), bias)
end)

This is part of the gradient checkpointing work in #1865. It replaces the list-form callback this PR first proposed, following the review suggestion to accept a container instead.

How it works

The block node keeps its arguments as the user wrote them, containers included. Everything that walks the arguments walks into the containers:

• The tracer converts each tensor inside a container to an expression, creates one body parameter per tensor in traversal order, and gives the user's function containers of parameters. Vectorized axes are tracked per tensor, so a container can mix vectorized and plain entries. The wrapper that re-vectorizes arguments counts a container as one argument.
• Nx.Defn.Tree.apply_args/4, the Evaluator, Nx.Defn.Graph and inspect traverse containers and leave keyword options alone.
• EXLA flattens the arguments into the compiled function's parameters in the same traversal order the tracer used. Native lowerings and EXLA.CustomCall see the arguments as written.
• Nx.block/4 and the Evaluator pick the backend from the tensors inside the containers.

Backends that run the default implementation with apply(fun, [struct | args]) need no change, since the containers pass straight through to the function. The gradient re-traces the body from the stored arguments, so it needs no change either.

Tests

Nx: a map argument alongside a tensor, a tuple with a vectorized entry, a gradient through a tuple argument, a tuple of 300 tensors forward and gradient, and a graph split whose boundary runs through a container argument. EXLA: a tuple with a vectorized entry forward and gradient, and a tuple of 300 tensors. Full Nx suite and EXLA block tests on host pass.

🤖 Generated with Claude Code

[polvalente]

I believe this makes for a confusing interface. Does block accept a container as one of the arguments? If not, then that's the proper fix, I believe