Merge pull request #1170 from leloykun:fc--fix-muon

PiperOrigin-RevId: 714411374

optax/contrib/_muon.py

@@ -169,7 +169,7 @@ def update_fn(updates, state, params=None):
     # Apply Newton-schulz orthogonalization.
     updates = jax.tree.map(
         lambda x: orthogonalize_via_newton_schulz(x, ns_coeffs_, ns_steps, eps),
-        updates,
+        mu_hat,
     )
     if adaptive:
       # Scale the orthogonalized updates by the dual norm of the original