Nx.LinAlg.svd gradient broken for batched inputs

[blasphemetheus]

Nx.LinAlg.svd gradient raises on batched inputs. svd_grad pattern-matches {m, n} = Nx.shape(input) which immediately fails when the input is 3D or higher.

x = Nx.tensor([
      [[4.0, 3.0], [6.0, 3.0]],
      [[2.0, 1.0], [5.0, 7.0]]
    ])                                  # {2, 2, 2}

Nx.Defn.grad(x, fn t ->
  {u, s, vt} = Nx.LinAlg.svd(t)
  Nx.add(Nx.add(Nx.sum(u), Nx.sum(s)), Nx.sum(vt))
end)
# ** (MatchError) no match of right hand side value: {2, 2, 2}
#    at nx/lib/nx/lin_alg/svd.ex:312

The whole svd_grad body is written assuming rank-2 input. Fixing this is more than just threading batch axes through Nx.dot — the eye_k/eye_m/eye_n construction, Nx.new_axis(s_sq, 1), and Nx.make_diagonal(s_input) calls all need to become batch-aware.

Same class of bug as #1740 (eigh), #1741 (triangular_solve), and the LU batched-grad bug.

[blasphemetheus]

Hmm this might better belong in meta issue just adding support for batch-awareness for these effected functions.

Will keep any other issues filed in a much more holistic package, less spammy