Fix Nx.LinAlg.pinv zero-branch shape for multiple batch dimensions

[blasphemetheus]

Nx.LinAlg.pinv raises for inputs with two or more batch dimensions of different sizes:

iex> Nx.LinAlg.pinv(Nx.iota({3, 2, 2, 2}, type: :f64) |> Nx.add(Nx.eye(2)))
** (ArgumentError) cannot broadcast tensor of dimensions {3, 2, 2, 2} to {2, 3, 2, 2}

pinv_zero_shape builds the shape of the all-zeros branch by reversing the dimension list, taking the first two, and rebuilding, but it never re-reverses the remaining batch prefix. With one batch dimension, or with equal batch dimensions, reversing is a no-op and nothing shows. With unequal ones the zero branch gets a batch-transposed shape, and cond raises when it unifies the two branch shapes.

The fix transposes the trailing two dimensions in place with put_elem and leaves the batch prefix alone.

The tests cover {3, 2, 2, 2} and the non-square {5, 4, 2, 3}, asserting both the output shape and the Moore-Penrose identity A·P·A = A per batch, plus the all-zeros branch on a double batch. They fail on main with the error above and pass with the fix. The full nx suite passes.

🤖 Generated with Claude Code

[blasphemetheus]

this one leads to a pretty straightforward argument error, but I couldn't say why the fix works as this is some arcane math to me, pinv. so if there's a better solution to this failing test, probably should go with that

[polvalente]

I think we just missed an Enum.reverse here to undo the reverse above! Code below is cleaner anyway :)