Fix all-NaN output for complex inputs with 'schatten' preconditioning in Muon

[shoemoney]

_schatten_first_newton_schulz_iteration computes the Gram matrix as a = x @ x.T (optax/contrib/_muon.py:233 on main). For complex x this is not the Hermitian Gram matrix, so the modified first Newton-Schulz step diverges and orthogonalize_via_newton_schulz(..., preconditioning='schatten') returns an all-NaN array. A muon(preconditioning='schatten') step on a complex-valued model therefore NaN-poisons the parameters on the first update. Complex support is a tested contract of this code path: test_newton_schulz asserts complex Newton-Schulz produces a unitary matrix, but only exercised the default preconditioning.

The siblings in the same function family already take the conjugate: _aol_first_newton_schulz_iteration (optax/contrib/_muon.py:217) and _base_newton_schulz_iteration (optax/contrib/_muon.py:247) both use a = x @ x.T.conj(). Schatten, added alongside aol in #1602, is the only one that omitted it. .conj() is a no-op for real dtypes, so real-input behavior is unchanged.

Fix: use x @ x.T.conj() in the schatten first iteration. Test: test_newton_schulz is now parameterized over 'frobenius', 'aol', and 'schatten', so its existing real-orthogonality and complex-unitarity assertions cover all three modified first steps.

Verification (CPU, jax 0.11.1): before the fix, the schatten complex case fails with an all-NaN Gram matrix; after, optax/contrib/_muon_test.py passes (34 tests, 18 subtests), and the schatten complex output is unitary to ~1e-7 while the real schatten path is numerically unchanged (max deviation from identity ~1e-8). 'spectral' is deliberately not added to the parameterization: it shares _base_newton_schulz_iteration, which already conjugates, so it has no instance of this bug.

[shoemoney]

CI note: Build and check types with pyrefly failure is pre-existing on main and unrelated to this PR's changed files. The 10 pyrefly errors are in linear_algebra.py/linesearch.py/_make_pert.py/_projections.py with no overlap with _muon.py/_muon_test.py and reproduce identically on base commit 0391582 (main HEAD). No code fix required from this PR; rebase or label will clear it.

[sylvesterkaczmarek]

This looks right to me. The Schatten first step now builds the Hermitian Gram matrix like the AOL/base paths, while .conj() is a no-op for real inputs. Parameterizing the existing orthogonality/unitarity test also covers the missing complex Schatten case. I don't see a blocker here.