calc trace(G.T @ X) instead of norm(G.T @ X)

optax/contrib/_muon.py

@@ -108,7 +108,7 @@ def update_fn(updates, state, params=None):
     updates, _ = jax.lax.scan(muon_iterator, updates, muon_coeffs)
     if adaptive:
       # Scale the orthogonalized updates by the dual norm of the original updates.
-      updates = jax.tree.map(lambda x, y: jnp.linalg.norm(x.T @ y) * y, mu_hat, updates)
+      updates = jax.tree.map(lambda x, y: jnp.linalg.trace(x.T @ y) * y, mu_hat, updates)
     mu = otu.tree_cast(mu, mu_dtype)
     return updates, MuonState(count=count_inc, mu=mu)
   return base.GradientTransformation(init_fn, update_fn)