Support vectorized gradients via boundary devectorization (#1533)

[blasphemetheus]

Closes #1533 (pending review).

Replaces the existing per-op vectorization handling in Nx.Defn.Grad with a boundary wrapper.

Approach

At the entry of Grad.transform, align heterogenous vectorized inputs to the union of vec axes via Nx.broadcast_vectors. This matches the implicit alignment the forward pass already performs and lets the grad recursion run in a homogeneous case. The seed and expression args are devectorized, so grad operates in fully devectorized space — vec axes become leading "batch" dimensions of the devectorized tensors.

A batch_count is threaded through update_grads and the per-op grad/5 clauses. unbroadcast and :reshape use it to preserve batch dims (which should not be summed out as if they were ordinary broadcast dims). Most ops just pass it through. In to_grad, each gradient is re-vectorized with the appropriate vec axes before returning to the caller, so the user's grad result carries the same vec shape as their input.

Behavior for heterogenous vec axes

Each gradient in a heterogenous-axes call carries all output vec axes (output-shape semantics). Each (foo:i, bar:j) element holds the per-instance partial derivative; nothing is collapsed. Callers that want a parameter-shaped gradient sum over the foreign axes themselves.

x = Nx.tensor([2.0, 5.0]) |> Nx.vectorize(:foo)
y = Nx.tensor([3.0, 4.0]) |> Nx.vectorize(:bar)

{grad_x, grad_y} =
  Nx.Defn.grad({x, y}, fn {x, y} -> Nx.sum(Nx.multiply(x, y)) end)

# grad_x and grad_y both have vec[foo: 2, bar: 3], with per-instance partials:
#   grad_x[foo:i, bar:j] = y[bar:j]
#   grad_y[foo:i, bar:j] = x[foo:i]

Linalg simplifications

Nx.LinAlg.Cholesky and Nx.LinAlg.QR grad rules are updated to use rank-independent dot axes (Nx.dot(a, [-2], b, [-2])) so they survive the batched inputs that arise from devectorization. Cholesky also picks up an inlined Nx.eye(l) (was the deprecated l |> Nx.shape() |> Nx.eye()) and a small batch_transpose helper for swapping the last two axes regardless of rank.

Test coverage

The test suite in nx/test/nx/defn/grad_test.exs adds a check_vectorized_grad helper that compares vectorized grad against per-element stacked grads, plus coverage for:

• Edge cases: vectorize/devectorize inside grad, rename, reshape-then-vectorize, non-vectorized input → vectorized output, multiple vec axes, second-order grad, value_and_grad, grad of non-vec target with vec capture, grad w.r.t. tuple of mixed inputs, large vectorized batches, and constant-grad ops (all/any/argmax/argmin).
• Heterogenous-vec inputs: aligned at the boundary, with concrete-value tests for multiply, add, sin(add), and x^2 * y asserting per-instance partials directly.
• Individual ops touched by the rework: exp, multiply, sum over a 2D inner axis, reshape, concatenate, window_sum, dot with a captured matrix, cumulative_sum, conv, while, mixed vec/non-vec, two stacked vec axes, composed sigmoid(x @ w + b).

mix test test/nx/defn/grad_test.exs → 268 tests, 0 failures, 0 skipped. Full mix test from nx/ → 1353 doctests, 1262 tests, 0 failures, 0 skipped.

Out of scope — separate issues

Three independent vectorized-grad bugs surfaced while implementing the boundary wrapper but are not addressed here. Each carries its original reproducer:

• Dot gradient fails with mixed vectorized axes #1706 — unbroadcast inside the dot gradient does not account for vec axes; the contracted form fails for any vectorized input
• Vectorized gradient fails for apply_vectorized ops (Cholesky, triangular_solve, cond) #1729 — apply_vectorized ops (Nx.LinAlg.cholesky, Nx.LinAlg.triangular_solve, Nx.cond/Nx.select) re-vectorize internally, which the boundary-wrapper grad cannot reconcile with its devectorized recursion
• Vectorized gradient fails when input and output vec axis of the same name have different sizes #1730 — gradient fails when input and output share a vec axis name with different sizes

[polvalente]

Let's improve Nx.Testing.assert_all_close to support vectorization? We can do it in a follow-up PR

[polvalente]

assert_equal as well

[polvalente]

We need to test (for a 3D input, vectorized with [:batch]):

1. output with different names and length: input |> Nx.devectorize(keep_names: false) |> Nx.vectorize([:a, :b])
2. output with same name, but "hidden" vectors appear: input |> Nx.devectorize(keep_names: false) |> Nx.vectorize([:a, :b]) |> Nx.sum() |> Nx.devectorize(keep_names: false) |> Nx.vectorize([:batch])
3. output with same length but different names

[polvalente]

PR generally looks great! I think the test cases I suggested are covered, but I'd like to confirm.
And we should also add a test on Nx.LinAlg.qr grad (hidden vectors, but it has custom grad) and Nx.LinAlg.eigh (hidden vectors, but relies on the autograd engine)

[polvalente]

Suggested change

	assert Nx.shape(Nx.devectorize(grad, keep_names: false)) == {4, 2, 3}
	assert Nx.to_flat_list(grad) == List.duplicate(1.0, 24)
	assert Nx.devectorize(grad, keep_names: false) == Nx.broadcast(1.0, {4, 2, 3})

[polvalente]

Suggested change

	assert Nx.shape(Nx.devectorize(grad, keep_names: false)) == {4, 2, 3}
	assert Nx.to_flat_list(grad) == List.duplicate(1.0, 24)
	assert Nx.devectorize(grad, keep_names: false)) == Nx.broadcast(1.0, {4, 2, 3})

[polvalente]

I think this is right because this isn't too different than having names [:batch, nil, nil] -> names [:a, :b, nil] -> sum(axes: [2]).

[polvalente]

I think you should use the helper function you added above here

[polvalente]

@josevalim Overall this one looks good to me. WDYT?

[blasphemetheus]

@polvalente — while stress-testing a fuzz branch against a squash-merged simulation of this PR, I found that cholesky_grad is still broken for batched (3D+) inputs. The qr_grad rewrite on this branch uses batch_axes/1 + explicit batch axes on Nx.dot, and that works end-to-end; cholesky_grad uses Nx.dot([-2], l, [-2]) (axis-relative indices only, no batch axes), which is correct for 2D but produces a 4D intermediate when the input has a leading batch dim.

Minimal reproducer, run on this PR branch:

x = Nx.tensor([[[4.0, 2.0], [2.0, 5.0]], [[9.0, 3.0], [3.0, 10.0]]])
Nx.Defn.grad(x, fn a -> Nx.sum(Nx.LinAlg.cholesky(a)) end)
# ** (ArgumentError) incompatible dimensions for a and b on triangular solve
#    at nx/lib/nx/lin_alg/cholesky.ex:...

With a differently-shaped batched input the failure surfaces one line earlier as cannot broadcast tensor of dimensions {2, 4, 4, 2} to {2, 4, 4} — that's the 4D-from-dot problem I mentioned, caught at the subsequent Nx.divide(num, den).

Fix looks like mirroring the qr_grad treatment — add batch_axes/1 in cholesky.ex, change the dot to Nx.dot(a, [-2], ba, b, [-2], ba), and adapt the Nx.eye(l) |> Nx.add(1) denominator to match shape when batched.

Happy to either:

1. Extend this PR with the Cholesky fix (I can push a commit), or
2. File a follow-up issue/PR after this one merges and fix in a separate change.

This finding is the same class as #1740/#1741/#1742/#1743 (eigh / triangular_solve / LU / SVD batched grads). These are all followups anyway

[blasphemetheus]

That stuff belongs in it's own PR on reflection, sort of a specific shape of input that maybe isn't often given.