fix tests

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

tests/test_transform/factory/test_correctness.py

@@ -26,11 +26,11 @@
 
 
 @pytest.mark.parametrize("type", ("hadamard", "random-hadamard"))
-@pytest.mark.parametrize("randomized", (True, False))
-def test_correctness_linear(type, randomized):
+@pytest.mark.parametrize("randomize", (True, False))
+def test_correctness_linear(type, randomize):
     size = (4, 8)
     module = torch.nn.Linear(*size, bias=True)
-    scheme = TransformScheme(type=type, randomized=randomized)
+    scheme = TransformScheme(type=type, randomize=randomize)
     factory = TransformFactory.from_scheme(scheme, name="")
 
     input_tfm = factory.create_transform(
@@ -55,8 +55,8 @@ def test_correctness_linear(type, randomized):
 
 
 @pytest.mark.parametrize("type", ("hadamard", "random-hadamard"))
-@pytest.mark.parametrize("randomized", (True, False))
-def test_correctness_model(type, randomized, model_apply, offload=False):
+@pytest.mark.parametrize("randomize", (True, False))
+def test_correctness_model(type, randomize, model_apply, offload=False):
     # load model
     model = model_apply[0]
     if offload:
@@ -71,7 +71,7 @@ def test_correctness_model(type, randomized, model_apply, offload=False):
     # apply transforms
     config = TransformConfig(
         config_groups={
-            "": TransformScheme(type=type, randomized=randomized, apply=model_apply[1])
+            "": TransformScheme(type=type, randomize=randomize, apply=model_apply[1])
         }
     )
     apply_transform_config(model, config)
@@ -84,6 +84,6 @@ def test_correctness_model(type, randomized, model_apply, offload=False):
 @requires_gpu
 @requires_accelerate()
 @pytest.mark.parametrize("type", ("hadamard", "random-hadamard"))
-@pytest.mark.parametrize("randomized", (True, False))
-def test_correctness_model_offload(type, randomized, model_apply):
-    test_correctness_model(type, randomized, model_apply, offload=True)
+@pytest.mark.parametrize("randomize", (True, False))
+def test_correctness_model_offload(type, randomize, model_apply):
+    test_correctness_model(type, randomize, model_apply, offload=True)

tests/test_transform/factory/test_memory.py

@@ -29,9 +29,9 @@
 
 
 @pytest.mark.parametrize("type", ("hadamard", "random-hadamard"))
-@pytest.mark.parametrize("randomized", (True, False))
+@pytest.mark.parametrize("randomize", (True, False))
 @pytest.mark.parametrize("requires_grad", (True, False))
-def test_memory_sharing(type, randomized, requires_grad, offload=False):
+def test_memory_sharing(type, randomize, requires_grad, offload=False):
     # load model (maybe with offloading)
     model = TransformableModel(2, 2, 4, 4, 8, 8)
     if offload:
@@ -42,7 +42,7 @@ def test_memory_sharing(type, randomized, requires_grad, offload=False):
         config_groups={
             "": TransformScheme(
                 type=type,
-                randomzied=randomized,
+                randomize=randomize,
                 requires_grad=requires_grad,
                 apply=[
                     TransformArgs(targets="Linear", location="input"),
@@ -84,9 +84,6 @@ def test_memory_sharing(type, randomized, requires_grad, offload=False):
 @requires_gpu
 @requires_accelerate()
 @pytest.mark.parametrize("type", ("hadamard", "random-hadamard"))
-@pytest.mark.parametrize("randomized", (True, False))
-def test_memory_sharing_offload(
-    type,
-    randomized,
-):
-    test_memory_sharing(type, randomized, requires_grad=False, offload=True)
+@pytest.mark.parametrize("randomize", (True, False))
+def test_memory_sharing_offload(type, randomize):
+    test_memory_sharing(type, randomize, requires_grad=False, offload=True)