Address review feedback: norm restoration, validation, tests

- Add norm restoration step after row-wise normalization to match
  reference NorMuon algorithm (paper Algorithm 1)
- Fix ns_coeffs validation readability (double-negative → positive)
- Tighten NorMuon vs Muon benchmark comparison (5x → 2x)
- Parameterize scale_by_normuon tests across matrix shapes

Author: bzantium

optax/contrib/_normuon.py

@@ -113,7 +113,7 @@ def init_fn(params):
           f'ns_coeffs must have shape (3,) or (n, 3), got {ns_coeffs_.shape}'
       )
     if ns_coeffs_.ndim == 2:
-      if not ns_coeffs_.shape[0] <= ns_steps:
+      if ns_coeffs_.shape[0] > ns_steps:
         raise ValueError(f'Not enough coeffs to perform {ns_steps} steps')
       ns_coeffs_ = ns_coeffs_[-ns_steps:]
 
@@ -159,7 +159,12 @@ def _update_nu(o, nu_prev):
 
     # Normalize and apply adaptive scaling.
     def _normalize(o, nu_new):
+      # Restore original norm after row-wise normalization (paper Algorithm 1).
+      o_norm = jnp.linalg.norm(o, ord='fro')
       o_hat = o / (jnp.sqrt(nu_new[..., None]) + eps)
+      o_hat_norm = jnp.linalg.norm(o_hat, ord='fro')
+      o_hat = o_hat * (o_norm / (o_hat_norm + eps))
+      # Apply adaptive scaling.
       m_n = math.prod(o.shape[-2:]) if o.ndim >= 2 else o.shape[-1]
       frob = jnp.linalg.norm(o_hat, ord='fro')
       scale = normuon_scale * jnp.sqrt(m_n) / (frob + eps)

optax/contrib/_normuon_benchmark_test.py

@@ -109,11 +109,11 @@ def test_normuon_vs_muon_convergence(self):
         'NorMuon did not converge.',
     )
 
-    # NorMuon should not be more than 5x worse than Muon.
+    # NorMuon should not be more than 2x worse than Muon.
     self.assertLess(
         normuon_losses[-1],
-        5.0 * muon_losses[-1],
-        'NorMuon final loss is more than 5x worse than Muon.',
+        2.0 * muon_losses[-1],
+        'NorMuon final loss is more than 2x worse than Muon.',
     )
 
   def test_normuon_no_side_effects(self):

optax/contrib/_normuon_test.py

@@ -73,18 +73,21 @@ def loss_fn(params):
     final_loss = loss_fn(params)
     self.assertLess(final_loss, initial_loss * 0.5)
 
-  def test_scale_by_normuon_direct(self):
+  @parameterized.product(
+      shape=[(6, 4), (8, 8), (3, 10)],
+  )
+  def test_scale_by_normuon_direct(self, shape):
     """Test scale_by_normuon directly on a 2D input."""
     key = jax.random.key(3)
-    params = jax.random.normal(key, (6, 4))
+    params = jax.random.normal(key, shape)
     opt = _normuon.scale_by_normuon()
     state = opt.init(params)
-    grad = jax.random.normal(key, (6, 4))
+    grad = jax.random.normal(key, shape)
     updates, new_state = opt.update(grad, state)
-    self.assertEqual(updates.shape, (6, 4))
+    self.assertEqual(updates.shape, shape)
     self.assertTrue(jnp.all(jnp.isfinite(updates)))
     # Check that nu state has the right shape (rows only).
-    self.assertEqual(new_state.nu.shape, (6,))
+    self.assertEqual(new_state.nu.shape, (shape[0],))
 
   @parameterized.named_parameters(
       ('small', 1e-7),