adding additional skips only for MI350 for LLVM cannot scalarize ERROR

Author: root

tests/pallas/gpu_attention_test.py

@@ -104,6 +104,8 @@ def test_mqa(
       kv_seq_len,
       return_residuals,
   ):
+    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
 
@@ -183,6 +185,8 @@ def test_gqa(
       kv_seq_len,
       return_residuals,
   ):
+    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
 

tests/pallas/gpu_ops_test.py

@@ -175,6 +175,9 @@ 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()]:
+      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:
       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)
@@ -292,7 +295,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 '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):
       self.skipTest("Skip on ROCm: tests/pallas/gpu_ops_test.py::FusedAttentionTest::test_fused_attention_bwd[4, 6, 7, 8, 9]")
 

tests/pallas/gpu_paged_attention_test.py

@@ -122,6 +122,9 @@ 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()]:
+      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}: