Fix wrong property access

- property access must not be in a form of function call
- fix wrong construction of prefill kv cache when mask_as_input is true

PiperOrigin-RevId: 755434493

Author: protobird-git

ai_edge_torch/generative/examples/gemma3/decoder.py

@@ -201,7 +201,7 @@ def create_sliding_mask(
         torch.zeros_like(sliding_mask_bool, dtype=torch.float),
         torch.full_like(
             sliding_mask_bool,
-            self.config.get_causal_mask_value(),
+            self.config.causal_mask_value,
             dtype=torch.float,
         ),
     )
@@ -219,7 +219,7 @@ def compose_mask(
       mask = torch.logical_and(mask, pixel_mask)
     else:
       mask = torch.logical_or(mask, pixel_mask)
-    mask = torch.where(mask, 0, self.config.get_causal_mask_value())
+    mask = torch.where(mask, 0, self.config.causal_mask_value)
     return mask
 
   def build_pixel_mask(self, image_indices: torch.Tensor):

ai_edge_torch/generative/layers/model_config.py

@@ -251,5 +251,5 @@ def block_config(self, idx: int) -> TransformerBlockConfig:
     return self.block_configs[idx]
 
   @property
-  def get_causal_mask_value(self) -> float:
+  def causal_mask_value(self) -> float:
     return self.block_config(0).attn_config.causal_mask_value

ai_edge_torch/generative/utilities/converter.py

@@ -243,15 +243,13 @@ def _export_helper(
 
   prefill_masks = None
   if flags.FLAGS.mask_as_input:
-    prefill_masks = [
-        _build_mask(
-            flags.FLAGS.prefill_seq_lens,
-            flags.FLAGS.kv_cache_max_len,
-            config.get_causal_mask_value(),
-        )
-    ]
-
-  if prefill_masks:
+    prefill_masks = _build_mask(
+        flags.FLAGS.prefill_seq_lens,
+        flags.FLAGS.kv_cache_max_len,
+        config.causal_mask_value,
+    )
+    if not isinstance(prefill_masks, list):
+      prefill_masks = [prefill_masks]
     assert len(prefill_masks) == len(prefill_seq_lens)
 
   decode_token = torch.tensor(
@@ -321,7 +319,7 @@ def _export_helper(
       #  torch.triu(mask, diagonal=decode_position).unsqueeze(0).unsqueeze(0)
       #
       sample_kwargs['mask'] = _build_mask(
-          1, flags.FLAGS.kv_cache_max_len, config.get_causal_mask_value()
+          1, flags.FLAGS.kv_cache_max_len, config.causal_mask_value
       )
     if lora is not None:
       sample_kwargs['lora'] = lora