fix is_device_rocm function call

root · gulsumgudukbay · commit bcc849838b0c · 2025-08-05T11:18:21.000-05:00
diff --git a/tests/debugging_primitives_test.py b/tests/debugging_primitives_test.py
@@ -806,7 +806,7 @@ def f2(x):
     self._assertLinesEqual(output(), "hello: 0\nhello: 1\nhello: 2\nhello: 3\n")
 
   def test_unordered_print_with_pjit(self):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
         self.skipTest("Skip on ROCm: tests/debugging_primitives_test.py::DebugPrintParallelTest::test_unordered_print_with_pjit")
     def f(x):
       debug_print("{}", x, ordered=False)
@@ -843,7 +843,7 @@ def f(x):
     self.assertEqual(output(), "[0 1 2 3 4 5 6 7]\n")
 
   def test_unordered_print_of_pjit_of_while(self):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
         self.skipTest("Skip on ROCm: tests/debugging_primitives_test.py::DebugPrintParallelTest::test_unordered_print_of_pjit_of_while")
     def f(x):
       def cond(carry):
diff --git a/tests/export_harnesses_multi_platform_test.py b/tests/export_harnesses_multi_platform_test.py
@@ -79,7 +79,7 @@ def setUp(self):
                "operator to work around missing XLA support for pair-reductions")
   )
   def test_prim(self, harness: test_harnesses.Harness):
-    if jtu.is_device_rocm and "multi_array_shape" in harness.fullname:
+    if jtu.is_device_rocm() and "multi_array_shape" in harness.fullname:
       self.skipTest("Skip on ROCm: test_prim_multi_array_shape")
 
     if "eigh_" in harness.fullname:
@@ -194,7 +194,7 @@ def test_all_gather(self, *, dtype):
     self.export_and_compare_to_native(f, x)
 
   def test_random_with_threefry_gpu_kernel_lowering(self):
-    if jtu.is_device_rocm and jtu.get_rocm_version() > (6, 5):
+    if jtu.is_device_rocm() and jtu.get_rocm_version() > (6, 5):
         self.skipTest("Skip on ROCm: test_random_with_threefry_gpu_kernel_lowering")
 
     # On GPU we use a custom call for threefry2x32
diff --git a/tests/ffi_test.py b/tests/ffi_test.py
@@ -173,7 +173,7 @@ def fun(x):
   @jtu.sample_product(shape=[(6, 5), (4, 5, 6)])
   @jtu.run_on_devices("gpu", "cpu")
   def test_ffi_call(self, shape):
-    if jtu.is_device_rocm and str(self).split()[0] == "test_ffi_call0":
+    if jtu.is_device_rocm() and str(self).split()[0] == "test_ffi_call0":
       self.skipTest("Skip on ROCm: test_ffi_call0")
 
     x = self.rng().randn(*shape).astype(np.float32)
@@ -189,7 +189,7 @@ def test_ffi_call(self, shape):
   )
   @jtu.run_on_devices("gpu", "cpu")
   def test_ffi_call_batching(self, shape, vmap_method):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
       self.skipTest("Skip on ROCm: test_ffi_call_batching")
 
     shape = (10,) + shape
@@ -276,7 +276,7 @@ def test_invalid_result_type(self):
 
   @jtu.run_on_devices("gpu", "cpu")
   def test_shard_map(self):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
         self.skipTest("Skip on ROCm: tests/ffi_test.py::FfiTest::test_shard_map")
     mesh = jtu.create_mesh((len(jax.devices()),), ("i",))
     x = self.rng().randn(8, 4, 5).astype(np.float32)
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
@@ -2599,7 +2599,7 @@ def testAssociativeScanOfBools(self):
   def testAssociativeScanSolvingRegressionTest(self, shape):
     # This test checks that the batching rule doesn't raise for a batch
     # sensitive function (solve).
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
       self.skipTest("Skip on ROCm: testAssociativeScanSolvingRegressionTest")
 
     ms = np.repeat(np.eye(2).reshape(1, 2, 2), shape, axis=0)
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
@@ -1562,10 +1562,10 @@ def testTrimZerosNotOneDArray(self):
   )
   @jax.default_matmul_precision("float32")
   def testPoly(self, a_shape, dtype, rank):
-    if jtu.is_device_rocm and a_shape == (12,) and dtype in ( np.int32,  np.int8 )  and rank == 2:
+    if jtu.is_device_rocm() and a_shape == (12,) and dtype in ( np.int32,  np.int8 )  and rank == 2:
         self.skipTest(f"Skip on ROCm: testPoly: a_shape == (12,) and dtype == {dtype} and rank == 2")
 
-    if jtu.is_device_rocm and a_shape == (6,) and dtype == np.float32 and rank == 2:
+    if jtu.is_device_rocm() and a_shape == (6,) and dtype == np.float32 and rank == 2:
         self.skipTest("Skip on ROCm: testPoly: a_shape == (6,) and dtype == numpy.float32 and rank == 2")
 
     if dtype in (np.float16, jnp.bfloat16, np.int16):
@@ -2585,7 +2585,7 @@ def testDiagFlat(self, shape, dtype, k):
     a2_shape=one_dim_array_shapes,
   )
   def testPolyMul(self, a1_shape, a2_shape, dtype):
-    if jtu.is_device_rocm and str(self).split()[0] == "testPolyMul1":
+    if jtu.is_device_rocm() and str(self).split()[0] == "testPolyMul1":
       self.skipTest("Skip on ROCm: testPolyMul")
 
     rng = jtu.rand_default(self.rng())
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
@@ -874,7 +874,7 @@ def testTensordot(self, lhs_shape, rhs_shape, axes, dtype):
   @jax.default_matmul_precision("float32")
   def testSVD(self, b, m, n, dtype, full_matrices, compute_uv, hermitian, algorithm):
     if algorithm is not None:
-      if jtu.is_device_rocm and algorithm == lax.linalg.SvdAlgorithm.JACOBI and dtype in {np.float32, np.complex64}:
+      if jtu.is_device_rocm() and algorithm == lax.linalg.SvdAlgorithm.JACOBI and dtype in {np.float32, np.complex64}:
         self.skipTest("Skip on ROCm: testSVD Jacobi tests")
       if hermitian:
         self.skipTest("Hermitian SVD doesn't support the algorithm parameter.")
@@ -1006,7 +1006,7 @@ def testNumpyQrModes(self, shape, dtype, mode):
   )
   @jax.default_matmul_precision("float32")
   def testQr(self, shape, dtype, full_matrices):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
       self.skipTest("Skip on ROCm: tests/linalg_test.py::NumpyLinalgTest::testQr")
 
     if (jtu.test_device_matches(["cuda"]) and
@@ -1084,7 +1084,7 @@ def testQrInvalidDtypeCPU(self, shape=(5, 6), dtype=np.float16):
     dtype=float_types + complex_types,
   )
   def testQrBatching(self, shape, dtype):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
       self.skipTest("Skip on ROCm: tests/linalg_test.py::NumpyLinalgTest::testQrBatching")
 
     rng = jtu.rand_default(self.rng())
diff --git a/tests/multiprocess_gpu_test.py b/tests/multiprocess_gpu_test.py
@@ -93,7 +93,7 @@ def test_gpu_distributed_initialize(self):
 
   def test_distributed_jax_visible_devices(self):
     """Test jax_visible_devices works in distributed settings."""
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
         self.skipTest("Skip on ROCm: test_distributed_jax_visible_devices")
     if not jtu.test_device_matches(['gpu']):
       raise unittest.SkipTest('Tests only for GPU.')
diff --git a/tests/nn_test.py b/tests/nn_test.py
@@ -201,7 +201,7 @@ def fwd(a, b, is_ref=False):
       impl=['cudnn', 'xla'],
   )
   def testDotProductAttention(self, dtype, group_num, use_vmap, impl):
-    if jtu.is_device_rocm and dtype == jnp.float16 and group_num == 4 and impl == 'xla':
+    if jtu.is_device_rocm() and dtype == jnp.float16 and group_num == 4 and impl == 'xla':
         self.skipTest("Skip on ROCm: testDotProductAttention[21,23]")
     if impl == 'cudnn' and not _is_required_cudnn_version_satisfied("8.0", 8904):
       raise unittest.SkipTest("CUDA or cuDNN versions are not compatible.")
diff --git a/tests/pallas/gpu_attention_test.py b/tests/pallas/gpu_attention_test.py
@@ -104,7 +104,7 @@ def test_mqa(
       kv_seq_len,
       return_residuals,
   ):
-    if jtu.is_device_rocm and 'gfx950' in [d.compute_capability for d in jax.devices()]:
+    if jtu.is_device_rocm() and 'gfx950' in [d.compute_capability for d in jax.devices()]:
       self.skipTest("Skip on ROCm: test_mqa: LLVM ERROR: Do not know how to scalarize the result of this operator!")
     del kwargs
     normalize_output = not return_residuals
@@ -185,7 +185,7 @@ def test_gqa(
       kv_seq_len,
       return_residuals,
   ):
-    if jtu.is_device_rocm and 'gfx950' in [d.compute_capability for d in jax.devices()]:
+    if jtu.is_device_rocm() and 'gfx950' in [d.compute_capability for d in jax.devices()]:
       self.skipTest("Skip on ROCm: test_gqa: LLVM ERROR: Do not know how to scalarize the result of this operator!")
     del kwargs
     normalize_output = not return_residuals
diff --git a/tests/pallas/gpu_ops_test.py b/tests/pallas/gpu_ops_test.py
@@ -175,10 +175,10 @@ def test_fused_attention_fwd(
       use_fwd,
       use_segment_ids,
   ):
-    if jtu.is_device_rocm and 'gfx950' in [d.compute_capability for d in jax.devices()]:
+    if jtu.is_device_rocm() and 'gfx950' in [d.compute_capability for d in jax.devices()]:
       self.skipTest("Skip on ROCm: test_fused_attention_fwd: LLVM ERROR: Do not know how to scalarize the result of this operator!")
 
-    if jtu.is_device_rocm and batch_size == 2 and seq_len == 384 and  num_heads == 8 and head_dim == 64 and block_sizes == (('block_q', 128), ('block_k', 128)) and causal and use_fwd and use_segment_ids:
+    if jtu.is_device_rocm() and batch_size == 2 and seq_len == 384 and  num_heads == 8 and head_dim == 64 and block_sizes == (('block_q', 128), ('block_k', 128)) and causal and use_fwd and use_segment_ids:
       self.skipTest("Skip on ROCm: tests/pallas/gpu_ops_test.py::FusedAttentionTest::test_fused_attention_fwd4")
     k1, k2, k3 = random.split(random.key(0), 3)
     q = random.normal(
@@ -232,7 +232,7 @@ def impl(q, k, v):
   # RESOURCE_EXHAUSTED: Shared memory size limit exceeded" error.
   @jtu.sample_product(
       batch_size=(1, 2),
-      seq_len=(32, 64) if jtu.is_device_rocm else (128, 384),
+      seq_len=(32, 64) if jtu.is_device_rocm() else (128, 384),
       num_heads=(1, 2),
       head_dim=(32, 64, 128,),
       block_sizes=(
@@ -253,7 +253,7 @@ def impl(q, k, v):
               ("block_kv_dq", 32),
           ),
       )
-      if jtu.is_device_rocm else (
+      if jtu.is_device_rocm() else (
           (
               ("block_q", 128),
               ("block_k", 128),
@@ -295,9 +295,9 @@ def test_fused_attention_bwd(
   ):
     test_name = str(self).split()[0]
     skip_suffix_list = [4, 6, 7, 8, 9]
-    if jtu.is_device_rocm and 'gfx950' in [d.compute_capability for d in jax.devices()]:
+    if jtu.is_device_rocm() and 'gfx950' in [d.compute_capability for d in jax.devices()]:
         self.skipTest("Skip on ROCm: test_fused_attention_bwd: LLVM ERROR: Do not know how to scalarize the result of this operator!")
-    if jtu.is_device_rocm and self.INTERPRET and any(test_name.endswith(str(suffix)) for suffix in skip_suffix_list):
+    if jtu.is_device_rocm() and self.INTERPRET and any(test_name.endswith(str(suffix)) for suffix in skip_suffix_list):
       self.skipTest("Skip on ROCm: tests/pallas/gpu_ops_test.py::FusedAttentionTest::test_fused_attention_bwd[4, 6, 7, 8, 9]")
 
     k1, k2, k3 = random.split(random.key(0), 3)
diff --git a/tests/pallas/gpu_paged_attention_test.py b/tests/pallas/gpu_paged_attention_test.py
@@ -122,12 +122,12 @@ def test_paged_attention(
       k_splits,
       attn_logits_soft_cap,
   ):
-    if jtu.is_device_rocm and 'gfx950' in [d.compute_capability for d in jax.devices()]:
+    if jtu.is_device_rocm() and 'gfx950' in [d.compute_capability for d in jax.devices()]:
       self.skipTest("Skip on ROCm: test_paged_attention: LLVM ERROR: Do not know how to scalarize the result of this operator!")
 
     test_name = str(self).split()[0]
     skip_numbers = {0, 1, 3, 5, 6, 7, 9}
-    if jtu.is_device_rocm and test_name in {f"test_paged_attention{i}" for i in skip_numbers}:
+    if jtu.is_device_rocm() and test_name in {f"test_paged_attention{i}" for i in skip_numbers}:
       self.skipTest("Skip on ROCm: tests/pallas/gpu_paged_attention_test.py::PagedAttentionKernelTest::test_paged_attention0")
     max_kv_len = 2048
     seq_lens = np.asarray([3, 256, 513, 1023, 2048], dtype=jnp.int32)
diff --git a/tests/pallas/ops_test.py b/tests/pallas/ops_test.py
@@ -528,7 +528,7 @@ def kernel(x_ref, o_ref):
   @hp.given(select_n_strategy(max_cases=2, min_rank=2, max_rank=4,
                               min_size_exp=1))
   def test_select_n(self, args):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
         self.skipTest("Skip on ROCm: test_select_n")
     pred, *cases = args
     scalar_pred = not pred.shape
@@ -560,7 +560,7 @@ def kernel(*refs):
   )
   @hp.given(hps.data())
   def test_unary_primitives(self, name, func, shape_dtype_strategy, data):
-    if jtu.is_device_rocm and name in {"logistic", "reciprocal"}:
+    if jtu.is_device_rocm() and name in {"logistic", "reciprocal"}:
         self.skipTest("Skip on ROCm: test_unary_primitives_[logistic,reciprocal]")
     if name in ["abs", "log1p", "pow2", "reciprocal", "relu", "sin", "sqrt"]:
       self.skip_if_mosaic_gpu()
@@ -1485,7 +1485,7 @@ def kernel(x_ref, y_ref, o_ref):
       for fn, dtype in itertools.product(*args)
   )
   def test_binary(self, f, dtype):
-    if jtu.is_device_rocm and f == jnp.bitwise_right_shift and dtype == "uint32":
+    if jtu.is_device_rocm() and f == jnp.bitwise_right_shift and dtype == "uint32":
         self.skipTest("Skip on ROCm: binary_bitwise_right_shift for uint32")
     self.skip_if_mosaic_gpu()
 
@@ -1559,7 +1559,7 @@ def kernel(o_ref):
 
   @parameterized.parameters("float16", "bfloat16", "float32")
   def test_approx_tanh(self, dtype):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
         self.skipTest("Skip on ROCm: test_approx_tanh")
 
     self.skip_if_mosaic_gpu()
@@ -1591,7 +1591,7 @@ def kernel(x_ref, o_ref):
     )
 
   def test_elementwise_inline_asm(self):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
         self.skipTest("Skip on ROCm: test_elementwise_inline_asm")
 
     self.skip_if_mosaic_gpu()
@@ -1876,7 +1876,7 @@ def f(x_ref, o_ref):
   def test_dot(self, lhs_and_rhs_shape, dtype, trans_x, trans_y):
     test_name = str(self).split()[0]
     skip_numbers = list(range(12, 20))
-    if jtu.is_device_rocm and jtu.get_rocm_version() == (7, 0) and test_name in {f"test_dot{i}" for i in skip_numbers}:
+    if jtu.is_device_rocm() and jtu.get_rocm_version() == (7, 0) and test_name in {f"test_dot{i}" for i in skip_numbers}:
         self.skipTest("Skip on ROCm: test_dot[12-19]")
     if (
         jtu.is_device_rocm() and
@@ -2169,7 +2169,7 @@ def masked_oob_swap_slice(_, _2, mask_ref, start_idx_ref, x_ref, y_ref):
       ("min_f32", pl.atomic_min, np.array([1, 2, 3, 4], np.float32), np.min),
   )
   def test_scalar_atomic(self, op, value, numpy_op):
-    if jtu.is_device_rocm and value.dtype == np.float32 and (op == pl.atomic_min or op == pl.atomic_max):
+    if jtu.is_device_rocm() and value.dtype == np.float32 and (op == pl.atomic_min or op == pl.atomic_max):
         self.skipTest("Skip on ROCm: test_scalar_atomic_(max/min)_f32")
 
     self.skip_if_mosaic_gpu()
@@ -2410,7 +2410,7 @@ def reduce(x_ref, y_ref):
       dtype=["float16", "float32", "int32", "uint32"],
   )
   def test_cumsum(self, dtype, axis):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
         self.skipTest("Skip on ROCm: test_cumsum")
     self.skip_if_mosaic_gpu()
 
diff --git a/tests/pallas/pallas_test.py b/tests/pallas/pallas_test.py
@@ -549,7 +549,7 @@ def index(x_ref, idx_ref, o_ref):
   def test_matmul(self, m, n, k, dtype, bm, bn, bk, gm):
     if jtu.test_device_matches(["tpu"]) and not self.INTERPRET:
       self.skipTest("On TPU the test works only in interpret mode")
-    if jtu.is_device_rocm and self.INTERPRET:
+    if jtu.is_device_rocm() and self.INTERPRET:
       self.skipTest("Skip on ROCm: test_matmul")
 
     k1, k2 = random.split(random.key(0))
diff --git a/tests/pallas/tpu_pallas_interpret_test.py b/tests/pallas/tpu_pallas_interpret_test.py
@@ -92,7 +92,7 @@ def setUp(self):
       self.skipTest(f'requires 1 device, found {self.num_devices}')
 
   def test_matmul_example(self):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
         self.skipTest("Skip on ROCm: test_matmul_example")
     def matmul_kernel(x_ref, y_ref, z_ref):
       z_ref[...] = x_ref[...] @ y_ref[...]
diff --git a/tests/pgle_test.py b/tests/pgle_test.py
@@ -481,7 +481,7 @@ def check_if_cache_hit(event):
   @parameterized.parameters([True, False])
   @jtu.thread_unsafe_test()
   def testAutoPgleWithCommandBuffers(self, enable_compilation_cache):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
         self.skipTest("Skip on ROCm: tests/pgle_test.py::PgleTest::testAutoPgleWithCommandBuffers")
     with (config.pgle_profiling_runs(1),
           config.enable_compilation_cache(enable_compilation_cache),
diff --git a/tests/python_callback_test.py b/tests/python_callback_test.py
@@ -948,7 +948,7 @@ def g(x):
     np.testing.assert_allclose(g(x), x)
 
   def test_can_shard_pure_callback_maximally(self):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
       self.skipTest("Skip on ROCm: test_can_shard_pure_callback_maximally")
     mesh = Mesh(np.array(jax.devices()), axis_names=('x',))
 
@@ -969,7 +969,7 @@ def f(x):
     )
 
   def test_can_shard_pure_callback_maximally_with_sharding(self):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
         self.skipTest("Skip on ROCm: test_can_shard_pure_callback_maximally_with_sharding")
     mesh = Mesh(np.array(jax.devices()), axis_names=('x',))
 
@@ -1244,7 +1244,7 @@ def f(x, y):
   def test_can_use_io_callback_in_pjit(
       self, *, ordered: bool, with_sharding: bool
   ):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
         self.skipTest("Skip on ROCm: test_can_use_io_callback_in_pjit")
     devices = jax.devices()
     mesh = jax.sharding.Mesh(np.array(devices), ['dev'])
@@ -1304,7 +1304,7 @@ def f(x):
       self.assertIn(f"{{maximal device={callback_device_index}}}", stablehlo_ir)
 
   def test_sequence_pjit_io_callback_ordered(self):
-    if jtu.is_device_rocm:
+    if jtu.is_device_rocm():
         self.skipTest("Skip on ROCm: test_sequence_pjit_io_callback_ordered")
     # A sequence of pairs of calls to pjit(io_callback(ordered=True)) with each
     # pair on a different device assignment.
diff --git a/tests/random_lax_test.py b/tests/random_lax_test.py
@@ -740,7 +740,7 @@ def testLogistic(self, dtype):
   )
   @jax.default_matmul_precision("float32")
   def testOrthogonal(self, n, shape, dtype, m):
-    if jtu.is_device_rocm and not (n == 0 or m == 0):
+    if jtu.is_device_rocm() and not (n == 0 or m == 0):
       self.skipTest("Skip on ROCm: testOrthogonal[1-3, 8-9]")
 
     if m is None:
diff --git a/tests/scipy_signal_test.py b/tests/scipy_signal_test.py
@@ -350,7 +350,7 @@ def osp_fun(x):
   def testWelchWithDefaultStepArgsAgainstNumpy(
       self, *, shape, dtype, nperseg, noverlap, use_nperseg, use_noverlap,
       use_window, timeaxis):
-    if jtu.is_device_rocm and (not use_nperseg or (use_nperseg and use_noverlap and use_window)):
+    if jtu.is_device_rocm() and (not use_nperseg or (use_nperseg and use_noverlap and use_window)):
       raise unittest.SkipTest("Skip on ROCm: testWelchWithDefaultStepArgsAgainstNumpy[1,2,3,7,8]")
 
     if tuple(shape) == (2, 3, 389, 5) and nperseg == 17 and noverlap == 13:
diff --git a/tests/shape_poly_test.py b/tests/shape_poly_test.py
@@ -3624,11 +3624,11 @@ def tearDown(self):
       #one_containing="",
   )
   def test_harness(self, harness: PolyHarness):
-    if jtu.is_device_rocm and harness.group_name in {"qr", "svd"} and harness.dtype in {np.float32, np.complex64}:
+    if jtu.is_device_rocm() and harness.group_name in {"qr", "svd"} and harness.dtype in {np.float32, np.complex64}:
         raise unittest.SkipTest("Skip for ROCm: shape_poly_test svd and qr for float32 and complex64.")
-    if jtu.is_device_rocm and harness.group_name == "vmap_eigh" and harness.dtype == np.complex64:
+    if jtu.is_device_rocm() and harness.group_name == "vmap_eigh" and harness.dtype == np.complex64:
         raise unittest.SkipTest("Skip for ROCm: shape_poly_test vmap_eigh for complex64.")
-    if jtu.is_device_rocm and "vmap_qr_multi_array" in harness.fullname and harness.dtype == np.float32:
+    if jtu.is_device_rocm() and "vmap_qr_multi_array" in harness.fullname and harness.dtype == np.float32:
         raise unittest.SkipTest("Skip for ROCm: shape_poly_test vmap_qr_multi_array for comple64 (note: when printed out, harness.dtype is np.float32.)")
 
     # We do not expect the associative scan error on TPUs
diff --git a/tests/sparse_test.py b/tests/sparse_test.py

Original file line number	Diff line number	Diff line change
`@@ -201,7 +201,7 @@ def fwd(a, b, is_ref=False):`
`201`	`201`	`impl=['cudnn', 'xla'],`
`202`	`202`	`)`
`203`	`203`	`def testDotProductAttention(self, dtype, group_num, use_vmap, impl):`
`204`		`- if jtu.is_device_rocm and dtype == jnp.float16 and group_num == 4 and impl == 'xla':`
	`204`	`+ if jtu.is_device_rocm() and dtype == jnp.float16 and group_num == 4 and impl == 'xla':`
`205`	`205`	`self.skipTest("Skip on ROCm: testDotProductAttention[21,23]")`
`206`	`206`	`if impl == 'cudnn' and not _is_required_cudnn_version_satisfied("8.0", 8904):`
`207`	`207`	`raise unittest.SkipTest("CUDA or cuDNN versions are not compatible.")`