Handle mask properly in gemma2

PiperOrigin-RevId: 716258803

Author: ai-edge-bot

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -144,12 +144,13 @@ def forward(
     attn_config = self.config.block_config(0).attn_config
     n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
     rope = rotary_pos_emb.build_rope(input_pos, n_elem, attn_config.rotary_base)
-    mask = [
-        self.get_attention_mask(
-            self.config.block_config(i).attn_config.attn_type, input_pos
-        )
-        for i in range(self.config.num_layers)
-    ]
+    if mask is None:
+      mask = [
+          self.get_attention_mask(
+              self.config.block_config(i).attn_config.attn_type, input_pos
+          )
+          for i in range(self.config.num_layers)
+      ]
 
     return self._forward_with_embeds(
         input_embeds, rope, mask, input_pos, kv_cache, export_config
@@ -159,7 +160,7 @@ def _forward_with_embeds(
       self,
       input_embeds: torch.Tensor,
       rope: Tuple[torch.Tensor, torch.Tensor],
-      mask: List[torch.Tensor],
+      mask: torch.Tensor | List[torch.Tensor],
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
       export_config: Optional[model_builder.ExportConfig] = None,
@@ -174,17 +175,10 @@ def _forward_with_embeds(
       input_embeds = input_embeds * self.config.embedding_scale
     x = input_embeds
     updated_kv_entries = []
-    mask_input = mask is not None
     for i, block in enumerate(self.transformer_blocks):
-      mask = (
-          mask
-          if mask_input
-          else self.get_attention_mask(
-              block.config.attn_config.attn_type, input_pos
-          )
-      )
+      mask_entry = mask[i] if isinstance(mask, list) else mask
       kv_entry = kv_cache.caches[i] if kv_cache else None
-      x, kv_entry = block(x, rope, mask[i], input_pos, kv_entry)
+      x, kv_entry = block(x, rope, mask_entry, input_pos, kv_entry)
       if kv_entry:
         updated_kv_entries.append(kv_entry)
     updated_kv_cache = kv_utils.KVCache(tuple(updated_kv_entries))

ai_edge_torch/generative/examples/paligemma/decoder2.py

@@ -86,7 +86,6 @@ def forward(
         embeds_len = input_embeds.shape[1]
         mask = torch.zeros(embeds_len, self.config.kv_cache_max)
         mask[:, embeds_len:] = float("-inf")
-        mask = [mask] * self.config.num_layers
 
     return self._forward_with_embeds(
         input_embeds, rope, mask, input_pos, kv_cache, export_config