Merge pull request #3628 from AI-Hypercomputer:chengnuojin-fix-splash

PiperOrigin-RevId: 900966264

Author: Google-ML-Automation

tests/utils/sharding_info/deepseek2-16b/tpu7x-8/slice_1/rule_ep-as-cp_ici_fsdp_parallelism=-1_ici_expert_parallelism=2/input_shardings.json

@@ -66,6 +66,18 @@
         "PartitionSpec": "P('fsdp', 'expert', None, None)"
       }
     },
+    {
+      "attention_op/arr: int8[1,4,4]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P(None, 'expert')"
+      }
+    },
+    {
+      "attention_op/arr: int32[2048]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P('expert',)"
+      }
+    },
     {
       "attention_op/query: bfloat16[96,16,2048,192]": {
         "logic_axes": "Unknown",
@@ -132,6 +144,24 @@
         "PartitionSpec": "P('fsdp', 'expert', None)"
       }
     },
+    {
+      "moe/w0_kernel: bfloat16[64,2048,1408]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P('expert', None, None)"
+      }
+    },
+    {
+      "moe/w1_kernel: bfloat16[64,2048,1408]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P('expert', None, None)"
+      }
+    },
+    {
+      "moe/wo_kernel: bfloat16[64,1408,2048]": {
+        "logic_axes": "Unknown",
+        "PartitionSpec": "P('expert', None, None)"
+      }
+    },
     {
       "linears/x: bfloat16[96,2048,2816]": {
         "logic_axes": "('activation_batch', 'activation_length', 'activation_mlp')",