[0.6.0-UT] Skipping aborted tests (HIP runtime issue) (#530)

Author: gulsumgudukbay

tests/image_test.py

@@ -160,6 +160,9 @@ def testResizeUp(self, dtype, image_shape, target_shape, method):
   )
   def testResizeGradients(self, dtype, image_shape, target_shape, method,
                            antialias):
+    if jtu.is_device_rocm():
+        self.skipTest("Skip on ROCm: testResizeGradients. Test aborts due to HIP runtime issue")
+
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: (rng(image_shape, dtype),)
     jax_fn = partial(image.resize, shape=target_shape, method=method,

tests/lax_numpy_reducers_test.py

@@ -671,6 +671,9 @@ def testNanStdGrad(self):
   @jax.numpy_dtype_promotion('standard')  # This test explicitly exercises mixed type promotion
   @jax.default_matmul_precision('float32')
   def testCov(self, shape, dtype, y_shape, y_dtype, rowvar, ddof, bias, fweights, aweights):
+    if jtu.is_device_rocm():
+        self.skipTest("Skip on ROCm: testCov. Test aborts due to HIP runtime issue")
+
     rng = jtu.rand_default(self.rng())
     wrng = jtu.rand_positive(self.rng())
     wdtype = np.real(dtype(0)).dtype

tests/lax_scipy_sparse_test.py

@@ -134,6 +134,8 @@ def args_maker():
     dtype=float_types + complex_types,
   )
   def test_cg_as_solve(self, shape, dtype):
+    if jtu.is_device_rocm():
+        self.skipTest("Skip on ROCm: test_cg_as_solve. Test aborts due to HIP runtime issue")
 
     rng = jtu.rand_default(self.rng())
     a = rng(shape, dtype)

tests/linalg_test.py

@@ -1479,6 +1479,8 @@ def testLuOnZeroMatrix(self, lu):
     dtype=float_types + complex_types,
   )
   def testLuGrad(self, shape, dtype):
+    if jtu.is_device_rocm():
+        self.skipTest("Skip on ROCm: testLuGrad. Test aborts due to HIP runtime issue")
     rng = jtu.rand_default(self.rng())
     a = rng(shape, dtype)
     lu = vmap(jsp.linalg.lu) if len(shape) > 2 else jsp.linalg.lu
@@ -1709,6 +1711,9 @@ def testTriangularSolveSingularBatched(self):
     dtype=int_types + float_types + complex_types
   )
   def testExpm(self, n, batch_size, dtype):
+    if jtu.is_device_rocm():
+        self.skipTest("Skip on ROCm: testExpm. Test aborts due to HIP runtime issue")
+
     if (jtu.test_device_matches(["cuda"]) and
         _is_required_cuda_version_satisfied(12000)):
       self.skipTest("Triggers a bug in cuda-12 b/287345077")
@@ -1861,6 +1866,8 @@ def sp_func(a):
     dtype=float_types + complex_types,
   )
   def testIssue2131(self, n, dtype):
+    if jtu.is_device_rocm():
+        self.skipTest("Skip on ROCm: testIssue2131. Test aborts due to HIP runtime issue")
     args_maker_zeros = lambda: [np.zeros((n, n), dtype)]
     osp_fun = lambda a: osp.linalg.expm(a)
     jsp_fun = lambda a: jsp.linalg.expm(a)
@@ -1896,6 +1903,8 @@ def args_maker():
     dtype=float_types + complex_types,
   )
   def testExpmFrechet(self, n, dtype):
+    if jtu.is_device_rocm():
+        self.skipTest("Skip on ROCm: testExpmFrechet. Test aborts due to HIP runtime issue")
     rng = jtu.rand_small(self.rng())
     if dtype == np.float64 or dtype == np.complex128:
       target_norms = [1.0e-2, 2.0e-1, 9.0e-01, 2.0, 3.0]
@@ -1934,6 +1943,9 @@ def args_maker():
     dtype=float_types + complex_types,
   )
   def testExpmGrad(self, n, dtype):
+    if jtu.is_device_rocm():
+        self.skipTest("Skip on ROCm: testExpmGrad. Test aborts due to HIP runtime issue")
+
     rng = jtu.rand_small(self.rng())
     a = rng((n, n), dtype)
     if dtype == np.float64 or dtype == np.complex128:

tests/qdwh_test.py

@@ -65,6 +65,8 @@ def _testPolarDecomposition(self, a, u, h, tol):
 
   def _testQdwh(self, a, dynamic_shape=None):
     """Computes the polar decomposition and tests its basic properties."""
+    if jtu.is_device_rocm():
+        self.skipTest("Skip on ROCm: testQdwh. Test aborts due to HIP runtime issue")
     eps = jnp.finfo(a.dtype).eps
     u, h, iters, conv = qdwh.qdwh(a, dynamic_shape=dynamic_shape)
     tol = 13 * eps