unskip selected test cases

Unskipped the following tests:
- lax_control_flow_test: testAssociativeScanSolvingRegressionTest
- scipy_signal_test: testWelchWithDefaultStepArgsAgainst
- sparse_test: cuSparseTest, SparseObjectTest
- random_lax_test: DistributionsTest::testOrthogonal{1,7,8,9}
- pallas_test: PallasCallInterpretTest::test_matmul (various shapes/dtypes)
- lax_numpy_test: testPolyMul1
- gpu_ops_test: FusedAttentionInterpretTest::test_fused_attention_bwd{4,6,7,8} (change the sample test data)
- export_harnesses_multi_platform_test: PrimitiveTest::test_prim_qr_multi_array_shape (various)
- ann_test: pass for whole file

Author: phambinhfin

tests/export_harnesses_multi_platform_test.py

@@ -79,8 +79,8 @@ 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:
-      self.skipTest("Skip on ROCm: test_prim_multi_array_shape")
+    # 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:
       self.skipTest("Eigenvalues are sorted and it is not correct to compare "
@@ -194,8 +194,8 @@ 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):
-        self.skipTest("Skip on ROCm: test_random_with_threefry_gpu_kernel_lowering")
+    # 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
     with config.threefry_gpu_kernel_lowering(True):

tests/lax_control_flow_test.py

@@ -2599,8 +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():
-      self.skipTest("Skip on ROCm: testAssociativeScanSolvingRegressionTest")
+
 
     ms = np.repeat(np.eye(2).reshape(1, 2, 2), shape, axis=0)
     vs = np.ones((shape, 2))

tests/lax_numpy_test.py

@@ -2585,8 +2585,8 @@ 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":
-      self.skipTest("Skip on ROCm: testPolyMul")
+    # if jtu.is_device_rocm() and str(self).split()[0] == "testPolyMul1":
+    #   self.skipTest("Skip on ROCm: testPolyMul")
 
     rng = jtu.rand_default(self.rng())
     np_fun = lambda arg1, arg2: np.polymul(arg1, arg2)

tests/pallas/gpu_ops_test.py

@@ -235,12 +235,12 @@ def impl(q, k, v):
       head_dim=(32, 64, 128,),
       block_sizes=(
           (
-              ("block_q", 64),
-              ("block_k", 64),
-              ("block_q_dkv", 16),
-              ("block_kv_dkv", 16),
-              ("block_q_dq", 16),
-              ("block_kv_dq", 64),
+              ("block_q", 128),
+              ("block_k", 128),
+              ("block_q_dkv", 128),
+              ("block_kv_dkv", 128),
+              ("block_q_dq", 128),
+              ("block_kv_dq", 128),
           ),
           (
               ("block_q", 32),
@@ -293,8 +293,8 @@ 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 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]")
+    # 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)
     q = random.normal(

tests/pallas/pallas_test.py

@@ -549,8 +549,9 @@ 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:
-      self.skipTest("Skip on ROCm: test_matmul")
+    # if jtu.is_device_rocm() and self.INTERPRET:
+    #   pass
+    #   # self.skipTest("Skip on ROCm: test_matmul")
 
     k1, k2 = random.split(random.key(0))
     x = random.normal(k1, (m, k), dtype=dtype)

tests/random_lax_test.py

@@ -740,8 +740,6 @@ 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):
-      self.skipTest("Skip on ROCm: testOrthogonal[1-3, 8-9]")
 
     if m is None:
       m = n

tests/scipy_signal_test.py

@@ -350,8 +350,8 @@ 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)):
-      raise unittest.SkipTest("Skip on ROCm: testWelchWithDefaultStepArgsAgainstNumpy[1,2,3,7,8]")
+    # 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:
       raise unittest.SkipTest("Test fails for these inputs")