use parametrize

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

tests/test_transform/factory/test_correctness.py

@@ -25,18 +25,12 @@
 from tests.testing_utils import requires_accelerate, requires_gpu
 
 
-def scheme_kwargs():
-    all_types = TransformFactory.registered_names()
-    base = [{"type": type} for type in all_types]
-    randomized = [{"type": type, "randomize": True} for type in all_types]
-    return base + randomized
-
-
-@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
-def test_correctness_linear(scheme_kwargs):
+@pytest.mark.parametrize("type", TransformFactory.registered_names())
+@pytest.mark.parametrize("randomized", (True, False))
+def test_correctness_linear(type, randomized):
     size = (4, 8)
     module = torch.nn.Linear(*size, bias=True)
-    scheme = TransformScheme(**scheme_kwargs)
+    scheme = TransformScheme(type=type, randomized=randomized)
     factory = TransformFactory.from_scheme(scheme, name="")
 
     input_tfm = factory.create_transform(
@@ -60,8 +54,9 @@ def test_correctness_linear(scheme_kwargs):
     assert torch.allclose(true_output, output, atol=1e-5, rtol=0.0)
 
 
-@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
-def test_correctness_model(scheme_kwargs, model_apply, offload=False):
+@pytest.mark.parametrize("type", TransformFactory.registered_names())
+@pytest.mark.parametrize("randomized", (True, False))
+def test_correctness_model(type, randomized, model_apply, offload=False):
     # load model
     model = model_apply[0]
     if offload:
@@ -76,10 +71,7 @@ def test_correctness_model(scheme_kwargs, model_apply, offload=False):
     # apply transforms
     config = TransformConfig(
         config_groups={
-            "": TransformScheme(
-                **scheme_kwargs,
-                apply=model_apply[1],
-            )
+            "": TransformScheme(type=type, randomized=randomized, apply=model_apply[1])
         }
     )
     apply_transform_config(model, config)
@@ -91,6 +83,7 @@ def test_correctness_model(scheme_kwargs, model_apply, offload=False):
 
 @requires_gpu
 @requires_accelerate()
-@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
-def test_correctness_model_offload(scheme_kwargs, model_apply):
-    test_correctness_model(scheme_kwargs, model_apply, offload=True)
+@pytest.mark.parametrize("type", TransformFactory.registered_names())
+@pytest.mark.parametrize("randomized", (True, False))
+def test_correctness_model_offload(type, randomized, model_apply):
+    test_correctness_model(type, randomized, model_apply, offload=True)

tests/test_transform/factory/test_memory.py

@@ -29,15 +29,10 @@
 from tests.testing_utils import requires_accelerate, requires_gpu
 
 
-def scheme_kwargs():
-    all_types = TransformFactory.registered_names()
-    base = [{"type": type} for type in all_types]
-    randomized = [{"type": type, "randomize": True} for type in all_types]
-    return base + randomized
-
-
-@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
-def test_memory_sharing(scheme_kwargs, offload=False):
+@pytest.mark.parametrize("type", TransformFactory.registered_names())
+@pytest.mark.parametrize("randomized", (True, False))
+@pytest.mark.parametrize("requires_grad", (True, False))
+def test_memory_sharing(type, randomized, requires_grad, offload=False):
     # load model (maybe with offloading)
     model = TransformableModel(2, 2, 4, 4, 8, 8)
     if offload:
@@ -47,7 +42,9 @@ def test_memory_sharing(scheme_kwargs, offload=False):
     config = TransformConfig(
         config_groups={
             "": TransformScheme(
-                **scheme_kwargs,
+                type=type,
+                randomzied=randomized,
+                requires_grad=requires_grad,
                 apply=[
                     TransformArgs(targets="Linear", location="input"),
                     TransformArgs(targets="Linear", location="output"),
@@ -87,12 +84,10 @@ def test_memory_sharing(scheme_kwargs, offload=False):
 
 @requires_gpu
 @requires_accelerate()
-@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
-def test_memory_sharing_offload(scheme_kwargs):
-    test_memory_sharing(scheme_kwargs, offload=True)
-
-
-@pytest.mark.parametrize("scheme_kwargs", scheme_kwargs())
-def test_memory_sharing_training(scheme_kwargs):
-    scheme_kwargs["requires_grad"] = True
-    test_memory_sharing(scheme_kwargs, offload=False)
+@pytest.mark.parametrize("type", TransformFactory.registered_names())
+@pytest.mark.parametrize("randomized", (True, False))
+def test_memory_sharing_offload(
+    type,
+    randomized,
+):
+    test_memory_sharing(type, randomized, requires_grad=False, offload=True)