update skip_if_rocm import

Peter Y. Yeh · Peter Y. Yeh · commit d4910871a751 · 2025-01-17T12:15:47.000-06:00
lint
diff --git a/test/dtypes/test_affine_quantized.py b/test/dtypes/test_affine_quantized.py
@@ -2,7 +2,6 @@
 import unittest
 
 import torch
-from test_utils import skip_if_rocm
 from torch.testing._internal import common_utils
 from torch.testing._internal.common_utils import (
     TestCase,
@@ -22,6 +21,7 @@
     TORCH_VERSION_AT_LEAST_2_5,
     TORCH_VERSION_AT_LEAST_2_6,
     is_sm_at_least_89,
+    skip_if_rocm,
 )
 
 
diff --git a/test/dtypes/test_floatx.py b/test/dtypes/test_floatx.py
@@ -2,7 +2,6 @@
 import unittest
 
 import torch
-from test_utils import skip_if_rocm
 from torch.testing._internal.common_utils import (
     TestCase,
     instantiate_parametrized_tests,
@@ -28,7 +27,7 @@
     fpx_weight_only,
     quantize_,
 )
-from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, is_fbcode
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, is_fbcode, skip_if_rocm
 
 _DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else [])
 _Floatx_DTYPES = [(3, 2), (2, 2)]
diff --git a/test/float8/test_base.py b/test/float8/test_base.py
@@ -18,14 +18,13 @@
     TORCH_VERSION_AT_LEAST_2_5,
     is_sm_at_least_89,
     is_sm_at_least_90,
+    skip_if_rocm,
 )
 
 if not TORCH_VERSION_AT_LEAST_2_5:
     pytest.skip("Unsupported PyTorch version", allow_module_level=True)
 
 
-from test_utils import skip_if_rocm
-
 from torchao.float8.config import (
     CastConfig,
     Float8LinearConfig,
diff --git a/test/hqq/test_hqq_affine.py b/test/hqq/test_hqq_affine.py
@@ -1,7 +1,6 @@
 import unittest
 
 import torch
-from test_utils import skip_if_rocm
 
 from torchao.quantization import (
     MappingType,
@@ -11,6 +10,7 @@
 )
 from torchao.utils import (
     TORCH_VERSION_AT_LEAST_2_3,
+    skip_if_rocm,
 )
 
 cuda_available = torch.cuda.is_available()
diff --git a/test/integration/test_integration.py b/test/integration/test_integration.py
@@ -80,6 +80,7 @@
     benchmark_model,
     is_fbcode,
     is_sm_at_least_90,
+    skip_if_rocm,
     unwrap_tensor_subclass,
 )
 
@@ -90,7 +91,6 @@
 except ModuleNotFoundError:
     has_gemlite = False
 
-from test_utils import skip_if_rocm
 
 logger = logging.getLogger("INFO")
 
diff --git a/test/kernel/test_galore_downproj.py b/test/kernel/test_galore_downproj.py
@@ -8,10 +8,10 @@
 
 import torch
 from galore_test_utils import make_data
-from test_utils import skip_if_rocm
 
 from torchao.prototype.galore.kernels.matmul import set_tuner_top_k as matmul_tuner_topk
 from torchao.prototype.galore.kernels.matmul import triton_mm_launcher
+from torchao.utils import skip_if_rocm
 
 torch.manual_seed(0)
 
diff --git a/test/prototype/test_awq.py b/test/prototype/test_awq.py
@@ -5,13 +5,15 @@
 import torch
 
 from torchao.quantization import quantize_
-from torchao.utils import TORCH_VERSION_AT_LEAST_2_3, TORCH_VERSION_AT_LEAST_2_5
+from torchao.utils import (
+    TORCH_VERSION_AT_LEAST_2_3,
+    TORCH_VERSION_AT_LEAST_2_5,
+    skip_if_rocm,
+)
 
 if TORCH_VERSION_AT_LEAST_2_3:
     from torchao.prototype.awq import AWQObservedLinear, awq_uintx, insert_awq_observer_
 
-from test_utils import skip_if_rocm
-
 
 class ToyLinearModel(torch.nn.Module):
     def __init__(self, m=512, n=256, k=128):
diff --git a/test/prototype/test_low_bit_optim.py b/test/prototype/test_low_bit_optim.py
@@ -30,6 +30,7 @@
     TORCH_VERSION_AT_LEAST_2_4,
     TORCH_VERSION_AT_LEAST_2_5,
     get_available_devices,
+    skip_if_rocm,
 )
 
 try:
@@ -42,7 +43,6 @@
 except ImportError:
     lpmm = None
 
-from test_utils import skip_if_rocm
 
 _DEVICES = get_available_devices()
 
diff --git a/test/prototype/test_splitk.py b/test/prototype/test_splitk.py
@@ -13,9 +13,8 @@
 except ImportError:
     triton_available = False
 
-from test_utils import skip_if_rocm
 
-from torchao.utils import skip_if_compute_capability_less_than
+from torchao.utils import skip_if_compute_capability_less_than, skip_if_rocm
 
 
 @unittest.skipIf(not triton_available, "Triton is required but not available")
diff --git a/test/quantization/test_galore_quant.py b/test/quantization/test_galore_quant.py
@@ -13,12 +13,12 @@
     dequantize_blockwise,
     quantize_blockwise,
 )
-from test_utils import skip_if_rocm
 
 from torchao.prototype.galore.kernels import (
     triton_dequant_blockwise,
     triton_quantize_blockwise,
 )
+from torchao.utils import skip_if_rocm
 
 SEED = 0
 torch.manual_seed(SEED)
diff --git a/test/quantization/test_marlin_qqq.py b/test/quantization/test_marlin_qqq.py
@@ -3,7 +3,6 @@
 
 import pytest
 import torch
-from test_utils import skip_if_rocm
 from torch import nn
 from torch.testing._internal.common_utils import TestCase, run_tests
 
@@ -20,7 +19,7 @@
     MappingType,
     choose_qparams_and_quantize_affine_qqq,
 )
-from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, is_fbcode
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, is_fbcode, skip_if_rocm
 
 
 @unittest.skipIf(
diff --git a/test/sparsity/test_marlin.py b/test/sparsity/test_marlin.py
@@ -2,7 +2,6 @@
 
 import pytest
 import torch
-from test_utils import skip_if_rocm
 from torch import nn
 from torch.testing._internal.common_utils import TestCase, run_tests
 
@@ -16,7 +15,7 @@
 )
 from torchao.sparsity.marlin import inject_24, pack_to_marlin_24, unpack_from_marlin_24
 from torchao.sparsity.sparse_api import apply_fake_sparsity
-from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, skip_if_rocm
 
 
 class SparseMarlin24(TestCase):
diff --git a/test/test_utils.py b/test/test_utils.py
@@ -1,40 +1,11 @@
-import functools
 import unittest
 from unittest.mock import patch
 
-import pytest
 import torch
 
 from torchao.utils import TorchAOBaseTensor, torch_version_at_least
 
 
-def skip_if_rocm(message=None):
-    """Decorator to skip tests on ROCm platform with custom message.
-
-    Args:
-        message (str, optional): Additional information about why the test is skipped.
-    """
-
-    def decorator(func):
-        @functools.wraps(func)
-        def wrapper(*args, **kwargs):
-            if torch.version.hip is not None:
-                skip_message = "Skipping the test in ROCm"
-                if message:
-                    skip_message += f": {message}"
-                pytest.skip(skip_message)
-            return func(*args, **kwargs)
-
-        return wrapper
-
-    # Handle both @skip_if_rocm and @skip_if_rocm() syntax
-    if callable(message):
-        func = message
-        message = None
-        return decorator(func)
-    return decorator
-
-
 class TestTorchVersionAtLeast(unittest.TestCase):
     def test_torch_version_at_least(self):
         test_cases = [
diff --git a/torchao/utils.py b/torchao/utils.py
@@ -7,6 +7,7 @@
 from math import gcd
 from typing import Any, Callable, Tuple
 
+import pytest
 import torch
 import torch.nn.utils.parametrize as parametrize
 
@@ -161,6 +162,33 @@ def wrapper(*args, **kwargs):
     return decorator
 
 
+def skip_if_rocm(message=None):
+    """Decorator to skip tests on ROCm platform with custom message.
+
+    Args:
+        message (str, optional): Additional information about why the test is skipped.
+    """
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            if torch.version.hip is not None:
+                skip_message = "Skipping the test in ROCm"
+                if message:
+                    skip_message += f": {message}"
+                pytest.skip(skip_message)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    # Handle both @skip_if_rocm and @skip_if_rocm() syntax
+    if callable(message):
+        func = message
+        message = None
+        return decorator(func)
+    return decorator
+
+
 def compute_max_diff(output: torch.Tensor, output_ref: torch.Tensor) -> torch.Tensor:
     return torch.mean(torch.abs(output - output_ref)) / torch.mean(
         torch.abs(output_ref)

Original file line number	Diff line number	Diff line change
`@@ -1,7 +1,6 @@`
`1`	`1`	`import unittest`
`2`	`2`
`3`	`3`	`import torch`
`4`		`-from test_utils import skip_if_rocm`
`5`	`4`
`6`	`5`	`from torchao.quantization import (`
`7`	`6`	`MappingType,`
`@@ -11,6 +10,7 @@`
`11`	`10`	`)`
`12`	`11`	`from torchao.utils import (`
`13`	`12`	`TORCH_VERSION_AT_LEAST_2_3,`
	`13`	`+ skip_if_rocm,`
`14`	`14`	`)`
`15`	`15`
`16`	`16`	`cuda_available = torch.cuda.is_available()`
Original file line number	Diff line number	Diff line change
`@@ -13,12 +13,12 @@`
`13`	`13`	`dequantize_blockwise,`
`14`	`14`	`quantize_blockwise,`
`15`	`15`	`)`
`16`		`-from test_utils import skip_if_rocm`
`17`	`16`
`18`	`17`	`from torchao.prototype.galore.kernels import (`
`19`	`18`	`triton_dequant_blockwise,`
`20`	`19`	`triton_quantize_blockwise,`
`21`	`20`	`)`
	`21`	`+from torchao.utils import skip_if_rocm`
`22`	`22`
`23`	`23`	`SEED = 0`
`24`	`24`	`torch.manual_seed(SEED)`