Addressed MR comments

Signed-off-by: mikail <mkhona@nvidia.com>

Author: mkhona-nvidia

emerging_optimizers/soap/moso.py

@@ -250,10 +250,9 @@ def _update_one_sided_momentum_factor(
     shampoo_beta: float,
 ) -> None:
     """Update the smaller-side covariance of the Muon momentum."""
-    if momentum.shape[0] <= momentum.shape[1]:
-        momentum_factor.lerp_(momentum @ momentum.T, 1 - shampoo_beta)
-    else:
-        momentum_factor.lerp_(momentum.T @ momentum, 1 - shampoo_beta)
+    left_preconditioned = momentum.shape[0] <= momentum.shape[1]
+    maybe_transposed_momentum = momentum if left_preconditioned else momentum.T
+    momentum_factor.lerp_(maybe_transposed_momentum @ maybe_transposed_momentum.T, 1 - shampoo_beta)
 
 
 @torch.no_grad()  # type: ignore[misc]
@@ -282,13 +281,13 @@ def _update_eigenbasis_and_adam_exp_avgs(
     )
 
     if use_eigh:
-        updated_eigenbasis = soap_utils.get_eigenbasis_eigh([momentum_factor])[0]
+        (updated_eigenbasis,) = soap_utils.get_eigenbasis_eigh([momentum_factor])
     else:
-        updated_eigenbasis = soap_utils.get_eigenbasis_qr(
+        (updated_eigenbasis,) = soap_utils.get_eigenbasis_qr(
             [momentum_factor],
             [eigenbasis],
             power_iter_steps=power_iter_steps,
-        )[0]
+        )
 
     exp_avg = _project_to_one_sided_eigenbasis(
         x=exp_avg,
@@ -307,7 +306,7 @@ def _sort_one_sided_eigenbasis_and_exp_avg_sq(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Sort eigenbasis slots by approximate eigenvalue and permute Adam second moments."""
     approx_eigvals = utils.eig.conjugate(momentum_factor, eigenbasis, diag=True)
-    sort_idx = torch.argsort(approx_eigvals, descending=True)
+    sort_idx = torch.argsort(approx_eigvals, descending=True, stable=True)
     sorted_eigenbasis = eigenbasis[:, sort_idx]
     exp_avg_sq_dim = 0 if left_preconditioned else 1
     return sorted_eigenbasis, exp_avg_sq.index_select(exp_avg_sq_dim, sort_idx)

tests/test_moso.py

@@ -93,7 +93,7 @@ def test_accumulates_momentum_covariance_on_smaller_side(self, shape: tuple[int,
         {"shape": (4, 8)},
         {"shape": (8, 4)},
     )
-    def test_no_ema_matches_one_sided_adam_in_eigenbasis(self, shape: tuple[int, int]) -> None:
+    def test_no_ema_is_close_to_one_sided_adam_in_eigenbasis(self, shape: tuple[int, int]) -> None:
         torch.manual_seed(7)
         grad = torch.randn(shape, device=FLAGS.device)
         param = torch.zeros(shape, requires_grad=True, device=FLAGS.device)
@@ -131,7 +131,7 @@ def test_no_ema_matches_one_sided_adam_in_eigenbasis(self, shape: tuple[int, int
         torch.testing.assert_close(
             applied_update,
             expected_update,
-            atol=1e-4,
+            atol=0.0,
             rtol=1e-4,
             msg=lambda msg: f"MOSO no-EMA update did not match projected Adam update for shape {shape}:\n{msg}",
         )