addressed some MR comments

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

Author: mkhona-nvidia

emerging_optimizers/scalar_optimizers/laprop.py

@@ -65,8 +65,7 @@ class LaProp(WeightDecayMixin, torch.optim.Optimizer):
         eps: Term added to the denominator for numerical stability.
         weight_decay: Weight decay coefficient.
         correct_bias: Whether to apply bias correction to the first and second moment EMAs.
-        normalize: Whether to normalize each updated parameter back to its pre-update Frobenius norm.
-        normalize_eps: Term used to avoid division by zero during parameter normalization.
+        frob_normalize: Whether to normalize each updated parameter back to its pre-update Frobenius norm.
         weight_decay_method: Method to apply weight decay, see
             :class:`~emerging_optimizers.mixin.WeightDecayMixin` for more details.
     """
@@ -79,8 +78,7 @@ def __init__(
         eps: float = 1e-8,
         weight_decay: float = 0.0,
         correct_bias: bool = True,
-        normalize: bool = False,
-        normalize_eps: float = 1e-8,
+        frob_normalize: bool = False,
         weight_decay_method: WeightDecayT = "decoupled",
     ) -> None:
         if not 0.0 <= lr:
@@ -91,16 +89,13 @@ def __init__(
             raise ValueError(f"Invalid beta at index 1: {betas[1]}")
         if not 0.0 <= eps:
             raise ValueError(f"Invalid epsilon value: {eps}")
-        if not 0.0 < normalize_eps:
-            raise ValueError(f"Invalid normalize_eps value: {normalize_eps}")
         if not 0.0 <= weight_decay:
             raise ValueError(f"Invalid weight_decay value: {weight_decay}")
-        if normalize and weight_decay != 0.0:
-            logging.warning("LaProp with normalize=True is intended to be used with weight_decay=0.0.")
+        if frob_normalize and weight_decay != 0.0:
+            logging.warning("LaProp with frob_normalize=True is intended to be used with weight_decay=0.0.")
         defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, correct_bias=correct_bias)
         self.weight_decay_method = weight_decay_method
-        self.normalize = normalize
-        self.normalize_eps = normalize_eps
+        self.frob_normalize = frob_normalize
         super().__init__(params, defaults)
 
     @torch.no_grad()
@@ -170,7 +165,7 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                 grad = p.grad
                 state = self.state[p]
                 state["step"] += 1
-                pre_norm = p.data.norm() if self.normalize else None
+                pre_norm = p.data.norm() if self.frob_normalize else None
 
                 self._apply_weight_decay_inplace(p.data, grad, lr, weight_decay)
 
@@ -184,7 +179,8 @@ def step(self, closure: Callable[[], float] | None = None) -> float | None:
                     eps,
                 )
                 p.data.add_(update, alpha=-lr)
-                if pre_norm is not None:
-                    p.data.mul_(pre_norm / p.data.norm().clamp_min(self.normalize_eps))
+                if self.frob_normalize:
+                    assert pre_norm is not None
+                    p.data.mul_(pre_norm / p.data.norm().clamp_min(eps))
 
         return loss

tests/test_laprop.py

@@ -71,8 +71,8 @@ def test_state_initialization(self, shape) -> None:
         expected_exp_avg_sq = (1 - beta2) * grad.square()
         normalized_grad = grad / (grad.abs() + optimizer.param_groups[0]["eps"])
         expected_exp_avg = (1 - beta1) * normalized_grad
-        torch.testing.assert_close(optimizer.state[param]["exp_avg_sq"], expected_exp_avg_sq)
-        torch.testing.assert_close(optimizer.state[param]["exp_avg"], expected_exp_avg)
+        torch.testing.assert_close(optimizer.state[param]["exp_avg_sq"], expected_exp_avg_sq, atol=0, rtol=0)
+        torch.testing.assert_close(optimizer.state[param]["exp_avg"], expected_exp_avg, atol=0, rtol=0)
 
     @parameterized.parameters(
         {"shape": (3, 3)},
@@ -96,9 +96,9 @@ def test_optimizer_step_matches_update_function(self, shape) -> None:
         param.grad = grad.clone()
         optimizer.step()
 
-        torch.testing.assert_close(param, old_param - lr * expected_update)
-        torch.testing.assert_close(optimizer.state[param]["exp_avg"], exp_avg)
-        torch.testing.assert_close(optimizer.state[param]["exp_avg_sq"], exp_avg_sq)
+        torch.testing.assert_close(param, old_param - lr * expected_update, atol=0, rtol=0)
+        torch.testing.assert_close(optimizer.state[param]["exp_avg"], exp_avg, atol=0, rtol=0)
+        torch.testing.assert_close(optimizer.state[param]["exp_avg_sq"], exp_avg_sq, atol=0, rtol=0)
 
     @parameterized.parameters(
         {"shape": (3, 3)},
@@ -118,10 +118,10 @@ def test_no_grad_no_update_params_unchanged(self, shape) -> None:
         {"shape": (15, 31)},
         {"shape": (127, 255)},
     )
-    def test_normalize_preserves_parameter_norm(self, shape) -> None:
+    def test_frob_normalize_preserves_parameter_norm(self, shape) -> None:
         """LaProp can normalize updated parameters back to their pre-update norm."""
         param = torch.nn.Parameter(torch.randint(1, 5, shape, device=self.device, dtype=torch.float32))
-        optimizer = LaProp([param], lr=0.25, weight_decay=0.0, normalize=True)
+        optimizer = LaProp([param], lr=0.25, weight_decay=0.0, frob_normalize=True)
         param.grad = torch.randint(1, 5, shape, device=self.device, dtype=torch.float32)
         original_norm = param.norm()
 
@@ -168,12 +168,6 @@ def test_negative_eps_raises_value_error(self) -> None:
         with self.assertRaisesRegex(ValueError, "Invalid epsilon"):
             LaProp([param], eps=-1e-8)
 
-    def test_non_positive_normalize_eps_raises_value_error(self) -> None:
-        """Test that LaProp raises ValueError for non-positive normalize_eps."""
-        param = torch.nn.Parameter(torch.randn(3, 3, device=self.device))
-        with self.assertRaisesRegex(ValueError, "Invalid normalize_eps"):
-            LaProp([param], normalize_eps=0.0)
-
     def test_negative_weight_decay_raises_value_error(self) -> None:
         """Test that LaProp raises ValueError for negative weight_decay."""
         param = torch.nn.Parameter(torch.randn(3, 3, device=self.device))