Refactor tests to use pytest.mark.parametrize instead of for loops

Author: voltjia

tests/test_div.py

@@ -7,19 +7,22 @@
 
 
 @skip_if_cuda_not_available
+@pytest.mark.parametrize(
+    "rounding_mode",
+    [
+        None,
+        pytest.param(
+            "trunc", marks=pytest.mark.skip(reason="TODO: Test for `trunc` mode later.")
+        ),
+        "floor",
+    ],
+)
 @pytest.mark.parametrize(*generate_arguments())
-def test_div(shape, dtype, device, rtol, atol):
+def test_div(shape, rounding_mode, dtype, device, rtol, atol):
     input = torch.randn(shape, dtype=dtype, device=device)
     other = torch.randn(shape, dtype=dtype, device=device)
 
-    for rounding_mode in (None, "trunc", "floor"):
-        # TODO: Test for `trunc` mode later.
-        if rounding_mode == "trunc":
-            continue
+    ninetoothed_output = ntops.torch.div(input, other, rounding_mode=rounding_mode)
+    reference_output = torch.div(input, other, rounding_mode=rounding_mode)
 
-        ninetoothed_output = ntops.torch.div(input, other, rounding_mode=rounding_mode)
-        reference_output = torch.div(input, other, rounding_mode=rounding_mode)
-
-        assert torch.allclose(
-            ninetoothed_output, reference_output, rtol=rtol, atol=atol
-        )
+    assert torch.allclose(ninetoothed_output, reference_output, rtol=rtol, atol=atol)

tests/test_gelu.py

@@ -8,14 +8,20 @@
 
 
 @skip_if_cuda_not_available
+@pytest.mark.parametrize(
+    "approximate",
+    (
+        "none",
+        pytest.param(
+            "tanh", marks=pytest.mark.skip(reason="TODO: Test for `tanh` mode later.")
+        ),
+    ),
+)
 @pytest.mark.parametrize(*generate_arguments())
-def test_gelu(shape, dtype, device, rtol, atol):
+def test_gelu(shape, approximate, dtype, device, rtol, atol):
     input = torch.randn(shape, dtype=dtype, device=device)
 
-    for approximate in ("none", "tanh"):
-        ninetoothed_output = ntops.torch.gelu(input)
-        reference_output = F.gelu(input)
+    ninetoothed_output = ntops.torch.gelu(input, approximate=approximate)
+    reference_output = F.gelu(input, approximate=approximate)
 
-        assert torch.allclose(
-            ninetoothed_output, reference_output, rtol=rtol, atol=atol
-        )
+    assert torch.allclose(ninetoothed_output, reference_output, rtol=rtol, atol=atol)

tests/test_relu.py

@@ -8,14 +8,12 @@
 
 
 @skip_if_cuda_not_available
+@pytest.mark.parametrize("inplace", (False, True))
 @pytest.mark.parametrize(*generate_arguments())
-def test_relu(shape, dtype, device, rtol, atol):
+def test_relu(shape, inplace, dtype, device, rtol, atol):
     input = torch.randn(shape, dtype=dtype, device=device)
 
-    for inplace in (False, True):
-        ninetoothed_output = ntops.torch.relu(input, inplace)
-        reference_output = F.relu(input, inplace)
+    ninetoothed_output = ntops.torch.relu(input, inplace)
+    reference_output = F.relu(input, inplace)
 
-        assert torch.allclose(
-            ninetoothed_output, reference_output, rtol=rtol, atol=atol
-        )
+    assert torch.allclose(ninetoothed_output, reference_output, rtol=rtol, atol=atol)