try to skip using cutensor in tests involving nan/inf

leofang · Mar 14, 2024 · a8d17c9 · a8d17c9
1 parent 0347782
commit a8d17c9
Showing 1 changed file with 34 additions and 20 deletions.
diff --git a/tests/cupy_tests/core_tests/test_ndarray_reduction.py b/tests/cupy_tests/core_tests/test_ndarray_reduction.py
@@ -1,5 +1,3 @@
-import unittest
-
 import numpy
 import pytest
 
@@ -12,7 +10,33 @@
 @testing.parameterize(*testing.product({
     'order': ('C', 'F'),
 }))
-class TestArrayReduction(unittest.TestCase):
+class TestArrayReduction:
+
+    @pytest.fixture(scope='class')
+    def exclude_cutensor(self):
+        # cuTENSOR seems to have issues in handling inf/nan in reduction-based
+        # routines, so we use this fixture to skip testing it
+        self.old_routine_accelerators = _acc.get_routine_accelerators()
+        self.old_reduction_accelerators = _acc.get_reduction_accelerators()
+
+        rot_acc = self.old_routine_accelerators.copy()
+        try:
+            rot_acc.remove(_acc.ACCELERATOR_CUTENSOR)
+        except ValueError:
+            pass
+        _acc.set_routine_accelerators(rot_acc)
+
+        red_acc = self.old_reduction_accelerators.copy()
+        try:
+            red_acc.remove(_acc.ACCELERATOR_CUTENSOR)
+        except ValueError:
+            pass
+        _acc.set_reduction_accelerators(red_acc)
+
+        yield
+
+        _acc.set_routine_accelerators(self.old_routine_accelerators)
+        _acc.set_reduction_accelerators(self.old_reduction_accelerators)
 
     @testing.for_all_dtypes()
     @testing.numpy_cupy_allclose(contiguous_check=False)
@@ -64,9 +88,7 @@ def test_max_multiple_axes_keepdims(self, xp, dtype):
 
     @testing.for_float_dtypes()
     @testing.numpy_cupy_allclose(contiguous_check=False)
-    def test_max_nan(self, xp, dtype):
-        if _acc.ACCELERATOR_CUTENSOR in _acc.get_routine_accelerators():
-            pytest.skip()
+    def test_max_nan(self, xp, dtype, exclude_cutensor):
         a = xp.array([float('nan'), 1, -1], dtype, order=self.order)
         return a.max()
 
@@ -84,7 +106,7 @@ def test_max_nan_imag(self, xp, dtype):
 
     @testing.for_float_dtypes()
     @testing.numpy_cupy_allclose(contiguous_check=False)
-    def test_max_inf(self, xp, dtype):
+    def test_max_inf(self, exclude_cutensor, xp, dtype):
         # cupy/cupy#8180
         a = xp.array([-float('inf'), -float('inf')], dtype, order=self.order)
         return a.max()
@@ -139,9 +161,7 @@ def test_min_multiple_axes_keepdims(self, xp, dtype):
 
     @testing.for_float_dtypes()
     @testing.numpy_cupy_allclose(contiguous_check=False)
-    def test_min_nan(self, xp, dtype):
-        if _acc.ACCELERATOR_CUTENSOR in _acc.get_routine_accelerators():
-            pytest.skip()
+    def test_min_nan(self, xp, dtype, exclude_cutensor):
         a = xp.array([float('nan'), 1, -1], dtype, order=self.order)
         return a.min()
 
@@ -159,7 +179,7 @@ def test_min_nan_imag(self, xp, dtype):
 
     @testing.for_float_dtypes()
     @testing.numpy_cupy_allclose(contiguous_check=False)
-    def test_min_inf(self, xp, dtype):
+    def test_min_inf(self, xp, dtype, exclude_cutensor):
         # cupy/cupy#8180
         a = xp.array([float('inf'), float('inf')], dtype, order=self.order)
         return a.min()
@@ -218,9 +238,7 @@ def test_ptp_multiple_axes_keepdims(self, xp, dtype):
 
     @testing.for_float_dtypes()
     @testing.numpy_cupy_allclose(contiguous_check=False)
-    def test_ptp_nan(self, xp, dtype):
-        if _acc.ACCELERATOR_CUTENSOR in _acc.get_routine_accelerators():
-            pytest.skip()
+    def test_ptp_nan(self, xp, dtype, exclude_cutensor):
         a = xp.array([float('nan'), 1, -1], dtype, order=self.order)
         return a.ptp()
 
@@ -268,9 +286,7 @@ def test_argmax_axis2(self, xp, dtype):
 
     @testing.for_float_dtypes()
     @testing.numpy_cupy_allclose(contiguous_check=False)
-    def test_argmax_nan(self, xp, dtype):
-        if _acc.ACCELERATOR_CUTENSOR in _acc.get_routine_accelerators():
-            pytest.skip()
+    def test_argmax_nan(self, xp, dtype, exclude_cutensor):
         a = xp.array([float('nan'), 1, -1], dtype, order=self.order)
         return a.argmax()
 
@@ -318,9 +334,7 @@ def test_argmin_axis2(self, xp, dtype):
 
     @testing.for_float_dtypes()
     @testing.numpy_cupy_allclose(contiguous_check=False)
-    def test_argmin_nan(self, xp, dtype):
-        if _acc.ACCELERATOR_CUTENSOR in _acc.get_routine_accelerators():
-            pytest.skip()
+    def test_argmin_nan(self, xp, dtype, exclude_cutensor):
         a = xp.array([float('nan'), 1, -1], dtype, order=self.order)
         return a.argmin()