fix(defn): accept constant tensor arguments in Nx.block

[cash-mckeeman]

Closes #1805.

Nx.block/4 raised whenever one of its arguments was a plain constant tensor rather than a traced expression. Nx.LinAlg.solve/2 is the only multi-argument block op today, so it is where this surfaced: solving against a closed-over right-hand side failed.

This is not gradient-specific — it reproduces under plain Nx.Defn.jit/1:

a = Nx.tensor([[4.0, 1.0, 0.5], [1.0, 3.0, 0.2], [0.5, 0.2, 2.0]])
b = Nx.tensor([1.0, 2.0, 3.0])

Nx.Defn.jit(fn x, y -> Nx.LinAlg.solve(x, y) end).(a, b)  # OK
Nx.Defn.jit(fn x -> Nx.LinAlg.solve(x, b) end).(a)        # ** (FunctionClauseError)

Present in released 0.13.0 and 0.13.1 as well as main.

Cause

expr_block/3 built one parameter per argument with parameter/2, which matches only %T{data: %Expr{context: context}} — it derives the tracing context from the argument itself, so a constant tensor had no matching clause.

Fix

Two parts, both needed:

1. Derive a fallback context from the first traced argument (block_context/1), and give it only to arguments that carry no context of their own. Arguments that were already expressions keep their own context exactly as before, so the change is confined to the constants that previously had no clause at all. If nothing is traced, the fallback is :root.
2. Convert the arguments to expressions and store those in the block node's in_args.

Part 2 is not cosmetic. With only part 1, parameter/2 stops raising but the block node still holds a raw tensor, and the failure moves to Nx.Defn.Tree.scope_ids_each/3, which likewise has no clause for a non-Expr tensor.

The per-argument contexts are captured before the to_expr/1 conversion, since converting erases the distinction between a traced argument and a constant.

The :root fallback in block_context/1 is defensive rather than reachable: Nx.block/4 dispatches through Nx.Shared.list_impl!(args), so an all-constant argument list resolves to BinaryBackend and never reaches expr_block/3. Happy to drop it if you would rather not carry unreachable branches.

Tests

Two, both verified failing on main before the fix:

• test/nx/defn_test.exs — a new describe "block", jitting a solve against a constant right-hand side. No test previously called Nx.block/4 directly, though the linalg grad suite drives the block path throughout.
• test/nx/defn/grad_test.exs — in the existing describe "solve", asserting the constant-rhs gradient equals the traced-rhs gradient computed by the existing solve_grad_wrt_a/2. Same mathematics, different tracing path, so any divergence is the defect.

Verification

• nx suite: 2725 passed, 0 failures (2723 before, plus these two).
• exla suite: 1416 passed, 0 failures, 49 excluded.
• mix compile --force: no new warnings.

Independent of #1806 / its PR — different files, both branched from main.

[polvalente]

I think this might have a slight bug, though I'm not sure if it's actually reproducible: If we have a block inside a while using a closure from outside the while, the order in which they appear might result in a hidden invalid context:

defn foo_lhs(a) do
  x = a + 1
  while {x, i = 0}, i < 10 do
    x_next = Nx.LinAlg.solve(a, x)
    {x, i + 1}
  end
end

defn foo_rhs(a) do
  x = a + 1
  while {x, i = 0}, i < 10 do
    x_next = Nx.LinAlg.solve(x, a)
    {x_next, i + 1}
  end
end

Neither of these should work, and I have a feeling that the second one will end up passing through.

[polvalente]

Both fail correctly in the branch, which is good. I am pushing a simplified implementation along with these as tests.

[polvalente]

I don't think it makes sense to have these tensors as module attributes

[polvalente]

We shouldn't have this as a module attribute. If you really need to reuse the value, either pass it down as a defn option or have it in the describe setup as a shared value.