unskip pallas tests that were skipped due to LLVm error

Author: gulsumgudukbay

tests/pallas/gpu_attention_test.py

@@ -104,8 +104,6 @@ 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
 
@@ -185,8 +183,6 @@ 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,8 +175,6 @@ 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")
@@ -295,8 +293,6 @@ 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,8 +122,6 @@ 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}