unblock splash kernel error using cp and explicit sharding

Author: NuojCheng

tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_1/rule_default/input_shardings.json

@@ -96,6 +96,18 @@
         "PartitionSpec": "P('fsdp', None, None, None)"
       }
     },
+    {
+      "attention_op/arr: int8[1,4,4]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P(None, None)"
+      }
+    },
+    {
+      "attention_op/arr: int32[2048]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P(None,)"
+      }
+    },
     {
       "attention_mla/out: bfloat16[192,2048,16,128]": {
         "logic_axes": "('activation_batch', 'activation_length', 'activation_heads', 'activation_kv')",

tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_1/rule_pure-fsdp/input_shardings.json

@@ -96,6 +96,18 @@
         "PartitionSpec": "P('fsdp', None, None, None)"
       }
     },
+    {
+      "attention_op/arr: int8[1,4,4]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P(None, None)"
+      }
+    },
+    {
+      "attention_op/arr: int32[2048]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P(None,)"
+      }
+    },
     {
       "attention_mla/out: bfloat16[192,2048,16,128]": {
         "logic_axes": "('activation_batch', 'activation_length', 'activation_heads', 'activation_kv')",

tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_1/rule_default/input_shardings.json

@@ -66,6 +66,24 @@
         "PartitionSpec": "P('fsdp', None, None, None)"
       }
     },
+    {
+      "attention_op/arr: int8[1,4,2]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P(None, None)"
+      }
+    },
+    {
+      "attention_op/arr: bool[2,512,512]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P()"
+      }
+    },
+    {
+      "attention_op/arr: int8[1,2,4]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P(None, None)"
+      }
+    },
     {
       "attentions/out: bfloat16[192,2048,64,64]": {
         "logic_axes": "('activation_batch', 'activation_attn_length', 'activation_heads', 'activation_kv')",

tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_1/rule_default_ici_fsdp_parallelism=-1_ici_expert_parallelism=2/input_shardings.json

@@ -66,6 +66,24 @@
         "PartitionSpec": "P(('fsdp', 'expert'), None, None, None)"
       }
     },
+    {
+      "attention_op/arr: int8[1,4,2]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P(None, None)"
+      }
+    },
+    {
+      "attention_op/arr: bool[2,512,512]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P()"
+      }
+    },
+    {
+      "attention_op/arr: int8[1,2,4]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P(None, None)"
+      }
+    },
     {
       "attentions/out: bfloat16[192,2048,64,64]": {
         "logic_axes": "('activation_batch', 'activation_attn_length', 'activation_heads', 'activation_kv')",

tests/utils/sharding_info/qwen3-0.6b/tpu7x-16/slice_1/rule_default/input_shardings.json

@@ -60,6 +60,18 @@
         "PartitionSpec": "P('fsdp', None, None, None)"
       }
     },
+    {
+      "attention_op/arr: int8[1,4,4]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P(None, None)"
+      }
+    },
+    {
+      "attention_op/arr: int32[2048]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P(None,)"
+      }
+    },
     {
       "attentions/out: bfloat16[192,2048,16,128]": {
         "logic_axes": "('activation_batch', 'activation_attn_length', 'activation_heads', 'activation_kv')",