cleaned up test cases

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

Author: mkhona-nvidia

tests/normalized_optimizer_convergence_test.py

@@ -60,8 +60,9 @@ def get_regular_parameters(self):
         return [self.fc3.weight, self.fc3.bias]
 
 
-# Base class for tests requiring seeding for determinism
-class BaseTestCase(parameterized.TestCase):
+class NormalizedOptimizerConvergenceTest(parameterized.TestCase):
+    """Convergence tests for normalized optimizers on a simple MLP task."""
+
     def setUp(self):
         """Set random seed before each test."""
         # Set seed for PyTorch
@@ -71,19 +72,17 @@ def setUp(self):
             torch.cuda.manual_seed_all(1234)
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-
-class NormalizedOptimizerConvergenceTest(BaseTestCase):
-    """Convergence tests for normalized optimizers on a simple MLP task."""
-
-    def _create_synthetic_mnist_data(self, num_samples=1000):
+    def _create_synthetic_mnist_data(self, num_samples: int = 1000) -> TensorDataset:
         """Create synthetic MNIST-like data for testing."""
         torch.manual_seed(1234)
         X = torch.randn(num_samples, 784, device=self.device)
         # Create somewhat realistic targets with class imbalance
         y = torch.randint(0, 10, (num_samples,))
         return TensorDataset(X, y)
 
-    def _train_model(self, model, optimizer_class, optimizer_kwargs, num_epochs=5):
+    def _train_model(
+        self, model: SimpleMLP, optimizer_class: torch.optim.Optimizer, optimizer_kwargs: dict, num_epochs: int = 5
+    ) -> tuple[float, float, float]:
         """Train model with given optimizer and return final loss and accuracy."""
         # Create data
         dataset = self._create_synthetic_mnist_data(num_samples=500)
@@ -140,7 +139,7 @@ def _train_model(self, model, optimizer_class, optimizer_kwargs, num_epochs=5):
 
         return initial_loss, final_loss, final_accuracy
 
-    def _verify_norms_preserved(self, model):
+    def _verify_norms_preserved(self, model: SimpleMLP) -> None:
         """Verify that oblique parameters maintain unit column norms."""
         for param in model.get_oblique_parameters():
             column_norms = param.data.norm(dim=0)  # Column norms
@@ -152,7 +151,7 @@ def _verify_norms_preserved(self, model):
                 rtol=1e-5,
             )
 
-    def test_oblique_sgd_convergence(self):
+    def test_oblique_sgd_convergence(self) -> None:
         """Test that ObliqueSGD can train a simple MLP and maintain norms."""
         model = SimpleMLP(input_size=784, hidden_size=64, num_classes=10)
 
@@ -168,7 +167,7 @@ def test_oblique_sgd_convergence(self):
         # Check norm preservation
         self._verify_norms_preserved(model)
 
-    def test_oblique_adam_convergence(self):
+    def test_oblique_adam_convergence(self) -> None:
         """Test that ObliqueAdam can train a simple MLP and maintain norms."""
         model = SimpleMLP(input_size=784, hidden_size=64, num_classes=10)
 
@@ -190,7 +189,7 @@ def test_oblique_adam_convergence(self):
         ("adam_col", ObliqueAdam, {"lr": 0.1, "betas": (0.9, 0.999), "weight_decay": 0.1, "dim": 0}),
         ("adam_row", ObliqueAdam, {"lr": 0.1, "betas": (0.9, 0.999), "weight_decay": 0.1, "dim": 1}),
     )
-    def test_optimizer_modes_convergence(self, optimizer_class, optimizer_kwargs):
+    def test_optimizer_modes_convergence(self, optimizer_class: torch.optim.Optimizer, optimizer_kwargs: dict) -> None:
         """Test that both row and column modes work for both optimizers."""
         model = SimpleMLP(input_size=784, hidden_size=32, num_classes=10)
 

tests/test_normalized_optimizer.py

@@ -19,8 +19,9 @@
 from emerging_optimizers.riemannian_optimizers.normalized_optimizer import ObliqueAdam, ObliqueSGD
 
 
-# Base class for tests requiring seeding for determinism
-class BaseTestCase(parameterized.TestCase):
+class NormalizedOptimizerFunctionalTest(parameterized.TestCase):
+    """Tests for ObliqueSGD and ObliqueAdam optimizers that preserve row/column norms."""
+
     def setUp(self):
         """Set random seed before each test."""
         # Set seed for PyTorch
@@ -30,15 +31,11 @@ def setUp(self):
             torch.cuda.manual_seed_all(1234)
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-
-class NormalizedOptimizerFunctionalTest(BaseTestCase):
-    """Tests for ObliqueSGD and ObliqueAdam optimizers that preserve row/column norms."""
-
     @parameterized.parameters(
         (0),
         (1),
     )
-    def test_oblique_sgd_preserves_norms(self, dim):
+    def test_oblique_sgd_preserves_norms(self, dim: int) -> None:
         """Test that ObliqueSGD preserves row or column norms after one optimization step."""
         # Create a 4x6 matrix for testing
         matrix_size = (4, 6)
@@ -76,7 +73,7 @@ def test_oblique_sgd_preserves_norms(self, dim):
         (0),
         (1),
     )
-    def test_oblique_adam_preserves_norms(self, dim):
+    def test_oblique_adam_preserves_norms(self, dim: int) -> None:
         """Test that ObliqueAdam preserves row or column norms after one optimization step."""
         # Create a 3x5 matrix for testing
         matrix_size = (3, 5)
@@ -109,7 +106,7 @@ def test_oblique_adam_preserves_norms(self, dim):
             rtol=1e-6,
         )
 
-    def test_oblique_sgd_zero_gradient(self):
+    def test_oblique_sgd_zero_gradient(self) -> None:
         """Test that ObliqueSGD handles zero gradients correctly."""
         matrix_size = (2, 4)
         param = torch.randn(matrix_size, dtype=torch.float32, device=self.device)
@@ -135,7 +132,7 @@ def test_oblique_sgd_zero_gradient(self):
         expected_norms = torch.ones_like(final_norms)
         torch.testing.assert_close(final_norms, expected_norms, atol=0, rtol=1e-6)
 
-    def test_oblique_adam_zero_gradient(self):
+    def test_oblique_adam_zero_gradient(self) -> None:
         """Test that ObliqueAdam handles zero gradients correctly."""
         matrix_size = (2, 3)
         param = torch.randn(matrix_size, dtype=torch.float32, device=self.device)
@@ -161,7 +158,7 @@ def test_oblique_adam_zero_gradient(self):
         expected_norms = torch.ones_like(final_norms)
         torch.testing.assert_close(final_norms, expected_norms, atol=0, rtol=1e-6)
 
-    def test_oblique_sgd_large_gradient(self):
+    def test_oblique_sgd_large_gradient(self) -> None:
         """Test that ObliqueSGD handles large gradients correctly."""
         matrix_size = (3, 4)
         param = torch.randn(matrix_size, dtype=torch.float32, device=self.device)
@@ -183,7 +180,7 @@ def test_oblique_sgd_large_gradient(self):
         expected_norms = torch.ones_like(final_norms)
         torch.testing.assert_close(final_norms, expected_norms, atol=0, rtol=1e-6)
 
-    def test_oblique_adam_large_gradient(self):
+    def test_oblique_adam_large_gradient(self) -> None:
         """Test that ObliqueAdam handles large gradients correctly."""
         matrix_size = (2, 5)
         param = torch.randn(matrix_size, dtype=torch.float32, device=self.device)
@@ -210,7 +207,7 @@ def test_oblique_adam_large_gradient(self):
             rtol=1e-6,
         )
 
-    def test_multiple_optimization_steps_preserve_norms(self):
+    def test_multiple_optimization_steps_preserve_norms(self) -> None:
         """Test that norms are preserved across multiple optimization steps."""
         matrix_size = (4, 4)
         param = torch.randn(matrix_size, dtype=torch.float32, device=self.device)
@@ -237,7 +234,7 @@ def test_multiple_optimization_steps_preserve_norms(self):
                 rtol=1e-6,
             )
 
-    def test_weight_decay_with_norm_preservation(self):
+    def test_weight_decay_with_norm_preservation(self) -> None:
         """Test that weight decay doesn't break norm preservation."""
         matrix_size = (3, 3)
         param = torch.randn(matrix_size, dtype=torch.float32, device=self.device)