Merge pull request #3505 from abhinavgoel95:patch-4

PiperOrigin-RevId: 889512992

Author: Google-ML-Automation

src/maxtext/configs/types.py

@@ -2282,12 +2282,17 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           "decoder_layer_input",
           "context",
           "mlpwi",
+          "moe_mlpwi_0",
+          "moe_mlpwi_1",
+          "moe_mlpwo",
           "mlpwi_0",
           "mlpwi_1",
           "mlpwo",
           "query_proj",
           "key_proj",
           "value_proj",
+          "query_wa_proj",
+          "kv_wa_proj",
           "mla_kv",
           "mla_q",
           "qkv_proj",