Fix Nx.LinAlg.pinv for batched input

[blasphemetheus]

Problem

Nx.LinAlg.pinv is broken for every batched (3-D+) input, with three size-dependent failure modes:

Input	Result
{2, 2, 2}	silent wrong output: returns shape {2, 2, 2, 2} (rank 4), no error
{2, n, n}, n ≥ 3	ArgumentError: cannot broadcast tensor of dimensions {2, n, 2, n} to {2, n, n}
non-square batched	same broadcast crash

Found by fuzzing; every other Nx.LinAlg decomposition supports leading batch dimensions.

Cause

The final contraction in pinv_non_zero used plain Nx.dot/2:

Nx.dot(v, sut)   # v: {b, n, k}, sut: {b, k, m}

For rank-3+ inputs Nx.dot/2 is a full tensor contraction, not a batched matmul — it produces a doubled-batch {b, n, b, m}. The defn cond guarding the all-zeros case then tries to unify the two branch shapes: for n ≥ 3 they can't broadcast (the crash), but for n = 2 the correct {2, 2, 2} zero-branch happens to broadcast into {2, 2, 2, 2} — which is why that case fails silently with a wrong-shape result.

Fix

One line, using the batch_axes/1 helper that already exists in lin_alg.ex for exactly this purpose:

batch_axes = batch_axes(v)
Nx.dot(v, [-1], batch_axes, sut, [-2], batch_axes)

Rank-2 input yields batch_axes == [] and identical behavior to before.

Tests

New "supports batched input" test in the pinv describe: {2,3,3}, {2,3,4}, {2,4,3}, and double-batch {2,2,3,3}, asserting (1) output shape, (2) the batched result equals the stacked per-matrix pinv results, and (3) the Moore–Penrose identity A·P·A ≈ A per batch. f64 inputs — the MP identity is conditioning-sensitive at f32 with unconditioned random draws.

Full nx suite green: 1384 doctests + 1370 tests.

Explicitly out of scope

Nx.LinAlg.svd itself crashes on any batched matrix with a size-1 dimension ({2,1,1}, {2,1,2}, {2,2,1} — unbatched {1,1} works), which independently blocks pinv for those shapes. That's a separate pre-existing bug and is not addressed here.

🤖 Generated with Claude Code