Leverage __subclasses__() to improve configs test (#72)

Summary:

Current tests on the configs with subclasses relied on explicit instantiating subclasses to test it. There are some limitations on this approach:
1. It is hard to catch newly added subclasses.
2. Due to some unknown interactions with `typing.Generic`, [the current `buck` test discovery mechanism is not able to discover the tests in it](#64 (comment)).

To resolve this, this diff refactors those tests with [`type.__subclasses()`](https://docs.python.org/3/reference/datamodel.html#type.__subclasses__) to tests the configs with subclasses.

Reviewed By: anana10c

Differential Revision: D67652761

distributed_shampoo/tests/shampoo_types_test.py

@@ -7,123 +7,59 @@
 
 """
 
+import itertools
 import re
 import unittest
-from abc import ABC, abstractmethod
-from typing import Generic, Type, TypeVar
 
 from distributed_shampoo.shampoo_types import (
     AdaGradGraftingConfig,
-    AdamGraftingConfig,
-    EigenvalueCorrectedShampooPreconditionerConfig,
     PreconditionerConfig,
     RMSpropGraftingConfig,
-    ShampooPreconditionerConfig,
 )
 
 
-class AdaGradGraftingConfigTest(unittest.TestCase):
+class AdaGradGraftingConfigSubclassesTest(unittest.TestCase):
     def test_illegal_epsilon(self) -> None:
         epsilon = 0.0
-        grafting_config_type = self._get_grafting_config_type()
-        with (
-            self.subTest(grafting_config_type=grafting_config_type),
-            self.assertRaisesRegex(
-                ValueError,
-                re.escape(f"Invalid epsilon value: {epsilon}. Must be > 0.0."),
-            ),
-        ):
-            grafting_config_type(epsilon=epsilon)
-
-    def _get_grafting_config_type(
-        self,
-    ) -> (
-        Type[AdaGradGraftingConfig]
-        | Type[RMSpropGraftingConfig]
-        | Type[AdamGraftingConfig]
-    ):
-        return AdaGradGraftingConfig
+        for cls in AdaGradGraftingConfig.__subclasses__():
+            with self.subTest(cls=cls):
+                self.assertRaisesRegex(
+                    ValueError,
+                    re.escape(f"Invalid epsilon value: {epsilon}. Must be > 0.0."),
+                    cls,
+                    epsilon=epsilon,
+                )
 
 
-class RMSpropGraftingConfigTest(AdaGradGraftingConfigTest):
+class RMSpropGraftingConfigSubclassesTest(AdaGradGraftingConfigSubclassesTest):
     def test_illegal_beta2(
         self,
     ) -> None:
-        grafting_config_type = self._get_grafting_config_type()
-        for beta2 in (-1.0, 0.0, 1.3):
-            with (
-                self.subTest(grafting_config_type=grafting_config_type, beta2=beta2),
+        for cls, beta2 in itertools.product(
+            RMSpropGraftingConfig.__subclasses__(), (-1.0, 0.0, 1.3)
+        ):
+            with self.subTest(cls=cls, beta2=beta2):
                 self.assertRaisesRegex(
                     ValueError,
                     re.escape(
                         f"Invalid grafting beta2 parameter: {beta2}. Must be in (0.0, 1.0]."
                     ),
-                ),
-            ):
-                grafting_config_type(beta2=beta2)
-
-    def _get_grafting_config_type(
-        self,
-    ) -> Type[RMSpropGraftingConfig] | Type[AdamGraftingConfig]:
-        return RMSpropGraftingConfig
-
-
-class AdamGraftingConfigTest(RMSpropGraftingConfigTest):
-    def _get_grafting_config_type(
-        self,
-    ) -> Type[RMSpropGraftingConfig] | Type[AdamGraftingConfig]:
-        return AdamGraftingConfig
-
-
-PreconditionerConfigType = TypeVar(
-    "PreconditionerConfigType", bound=PreconditionerConfig
-)
+                    cls,
+                    beta2=beta2,
+                )
 
 
-class AbstractPreconditionerConfigTest:
-    class PreconditionerConfigTest(
-        ABC,
-        unittest.TestCase,
-        Generic[PreconditionerConfigType],
-    ):
-        def test_illegal_num_tolerated_failed_amortized_computations(self) -> None:
-            num_tolerated_failed_amortized_computations = -1
-            with (
+class PreconditionerConfigSubclassesTest(unittest.TestCase):
+    def test_illegal_num_tolerated_failed_amortized_computations(self) -> None:
+        num_tolerated_failed_amortized_computations = -1
+        for cls in PreconditionerConfig.__subclasses__():
+            with self.subTest(cls=cls):
                 self.assertRaisesRegex(
                     ValueError,
                     re.escape(
                         f"Invalid num_tolerated_failed_amortized_computations value: "
                         f"{num_tolerated_failed_amortized_computations}. Must be >= 0."
                     ),
-                ),
-            ):
-                self._get_preconditioner_config_type()(
+                    cls,
                     num_tolerated_failed_amortized_computations=num_tolerated_failed_amortized_computations,
                 )
-
-        @abstractmethod
-        def _get_preconditioner_config_type(
-            self,
-        ) -> Type[PreconditionerConfigType]: ...
-
-
-class ShampooPreconditionerConfigTest(
-    AbstractPreconditionerConfigTest.PreconditionerConfigTest[
-        ShampooPreconditionerConfig
-    ]
-):
-    def _get_preconditioner_config_type(
-        self,
-    ) -> Type[ShampooPreconditionerConfig]:
-        return ShampooPreconditionerConfig
-
-
-class EigenvalueCorrectedShampooPreconditionerConfigTest(
-    AbstractPreconditionerConfigTest.PreconditionerConfigTest[
-        EigenvalueCorrectedShampooPreconditionerConfig
-    ]
-):
-    def _get_preconditioner_config_type(
-        self,
-    ) -> Type[EigenvalueCorrectedShampooPreconditionerConfig]:
-        return EigenvalueCorrectedShampooPreconditionerConfig