Add new code owner for multimodal workgroup

Author: hengtaoguo

.github/CODEOWNERS

@@ -1,2 +1,2 @@
 # Changes in this file should match with requiredReviewers in file .github/workflows/AddLabel.yml
-* @gobbleturk @khatwanimohit @bvandermoon @vipannalla @RissyRan @richjames0 @rni418 @gagika @shralex @yangyuwei @SurbhiJainUSC @hengtaoguo @A9isha @wang2yn84 @wyzhang @mitalisi @gpolovets1 @mailvijayasingh @jrplatin @patemotter @lumosis
+* @gobbleturk @khatwanimohit @bvandermoon @vipannalla @RissyRan @richjames0 @rni418 @gagika @shralex @yangyuwei @SurbhiJainUSC @hengtaoguo @A9isha @wang2yn84 @wyzhang @mitalisi @gpolovets1 @mailvijayasingh @jrplatin @patemotter @lumosis @aireenmei

.github/workflows/AddLabel.yml

@@ -74,6 +74,7 @@ jobs:
               jrplatin: "",
               patemotter: "",
               lumosis: "",
+              aireenmei: "",
             }
             const reviews = await github.rest.pulls.listReviews({
               owner,

MaxText/layers/llama4.py

@@ -428,6 +428,7 @@ def __call__(
     else:
       return layer_output
 
+
 class Llama4ScannableBlock(nn.Module):
   '''
   A repeatable block given nope_layer_interval and interleave_moe_layer_step
@@ -470,12 +471,12 @@ def __call__(
       nope_layer = determine_is_nope_layer(layer_id, self.nope_layer_interval)
       moe_layer = determine_is_moe_layer(layer_id, self.interleave_moe_layer_step)
       layer = Llama4DecoderLayer(
-        config=cfg,
-        mesh=mesh,
-        name=f"layers_{layer_id}",
-        quant=self.quant,
-        is_nope_layer=nope_layer,
-        is_moe_layer=moe_layer,
+          config=cfg,
+          mesh=mesh,
+          name=f"layers_{layer_id}",
+          quant=self.quant,
+          is_nope_layer=nope_layer,
+          is_moe_layer=moe_layer,
       )
       y = layer(
           y,
@@ -488,9 +489,8 @@ def __call__(
           slot=slot,
       )
       if cfg.scan_layers:
-        y=y[0]
+        y = y[0]
     if cfg.scan_layers:
       return y, None
     else:
       return y
-

MaxText/layers/models.py

@@ -362,6 +362,7 @@ def get_decoder_layers(self):
       return [simple_layer.SimpleMlpDecoderLayer]
     elif self.config.decoder_block == DecoderBlockType.LLAMA4:
       from MaxText.layers import llama4  # pylint: disable=import-outside-toplevel
+
       if self.config.scan_layers:
         return [llama4.Llama4ScannableBlock]
       else:
@@ -544,8 +545,8 @@ def __call__(
             layer_call_kwargs = {"bidirectional_mask": bidirectional_mask}
           elif cfg.decoder_block == DecoderBlockType.LLAMA4:
             layer_kwargs = {
-              "nope_layer_interval": self.config.nope_layer_interval,
-              "interleave_moe_layer_step": self.config.interleave_moe_layer_step,
+                "nope_layer_interval": self.config.nope_layer_interval,
+                "interleave_moe_layer_step": self.config.interleave_moe_layer_step,
             }
           RemattedBlockLayer = RemattedBlockLayers[0]
           scan_length = int(cfg.num_decoder_layers / cfg.inhomogeneous_layer_cycle_interval)

MaxText/max_utils.py

@@ -152,10 +152,9 @@ def maybe_initialize_jax_distributed_system(raw_keys):
     max_logging.log("Attempting to initialize the jax distributed system for CPU backend...")
     initialize_jax_for_cpu(raw_keys)
     max_logging.log("Jax distributed system initialized on CPUs!")
-  elif (
-      raw_keys["enable_checkpointing"]
-      and raw_keys["compile_topology_num_slices"] == -1
-  ) or raw_keys["hardware"] == "gpu_multiprocess":
+  elif (raw_keys["enable_checkpointing"] and raw_keys["compile_topology_num_slices"] == -1) or raw_keys[
+      "hardware"
+  ] == "gpu_multiprocess":
     max_logging.log("Attempting to initialize the jax distributed system...")
     if not raw_keys["enable_emergency_checkpoint"]:
       jax.distributed.initialize(initialization_timeout=raw_keys["jax_distributed_initialization_timeout"])

MaxText/tests/grpo_trainer_correctness_test.py

@@ -135,7 +135,7 @@ def setUp(self):
     )
     self.tokenizer_model.add_special_tokens({"pad_token": "<pad>"})
 
-  @pytest.mark.tpu_only # ATTENTION: Only run on TPU V4-8
+  @pytest.mark.tpu_only  # ATTENTION: Only run on TPU V4-8
   def test_grpo_trainer_correctness(self):
     # Get the expected (golden) data.
     golden_data = get_golden_data(self.config)