ai-edge-torch-nightly 0.3.0.dev20250116__py3-none-any.whl → 0.3.0.dev20250117__py3-none-any.whl

ai_edge_torch/_convert/conversion.py

@@ -109,21 +109,21 @@ def convert_signatures(
 
   _warn_training_modules(signatures)
 
-  def export(*args, **kwargs):
+  def export(**kwargs):
     nonlocal strict_export
     if strict_export == "auto":
       try:
-        exported_program = torch.export.export(*args, **kwargs, strict=True)
+        exported_program = torch.export.export(**kwargs, strict=True)
       except Exception:
         logging.warning(
             "torch.export.export(..., strict=True) failed. Retrying with"
             " strict=False"
         )
-        exported_program = torch.export.export(*args, **kwargs, strict=False)
+        exported_program = torch.export.export(**kwargs, strict=False)
     elif not strict_export:
-      exported_program = torch.export.export(*args, **kwargs, strict=False)
+      exported_program = torch.export.export(**kwargs, strict=False)
     else:
-      exported_program = torch.export.export(*args, **kwargs, strict=True)
+      exported_program = torch.export.export(**kwargs, strict=True)
 
     if hasattr(torch._decomp, "_decomp_table_to_post_autograd_aten"):
       # Available after torch 2.5.0: `_decomp_table_to_post_autograd_aten` is a
@@ -136,7 +136,12 @@ def convert_signatures(
     return exported_program
 
   exported_programs: torch.export.ExportedProgram = [
-      export(sig.module, sig.flat_args, dynamic_shapes=sig.dynamic_shapes)
+      export(
+          mod=sig.module,
+          args=sig.args,
+          kwargs=sig.kwargs,
+          dynamic_shapes=sig.dynamic_shapes,
+      )
       for sig in signatures
   ]
 

ai_edge_torch/_convert/signature.py

@@ -25,9 +25,9 @@ import torch.utils._pytree as pytree
 class Signature:
   name: str
   module: torch.nn.Module
-  sample_args: tuple[torch.Tensor]
+  sample_args: tuple[torch.Tensor, ...]
   sample_kwargs: dict[str, torch.Tensor]
-  dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None
+  dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any, ...]]] = None
 
   @property
   def _normalized_sample_args_kwargs(self):
@@ -61,6 +61,16 @@ class Signature:
     return names
 
   @property
-  def flat_args(self) -> tuple[Any]:
+  def flat_args(self) -> tuple[Any, ...]:
     args, kwargs = self._normalized_sample_args_kwargs
     return tuple([*args, *kwargs.values()])
+
+  @property
+  def args(self) -> tuple[Any, ...]:
+    args, _ = self._normalized_sample_args_kwargs
+    return args
+
+  @property
+  def kwargs(self) -> dict[str, Any]:
+    _, kwargs = self._normalized_sample_args_kwargs
+    return kwargs

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -144,12 +144,13 @@ class Gemma2(nn.Module):
     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 @@ class Gemma2(nn.Module):
       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 @@ class Gemma2(nn.Module):
       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 @@ class Decoder2(gemma2.Gemma2):
         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

ai_edge_torch/generative/examples/paligemma/image_encoder.py

@@ -135,7 +135,9 @@ def get_image_encoder_config() -> cfg.ModelConfig:
       image_embedding=image_embedding_config,
       block_configs=block_config,
       final_norm_config=norm_config,
-      enable_hlfb=True,
+      # TODO: b/377051577 - Once RemoveSDPACompositeZeroMaskPass is removed,
+      # enable_hlfb can be set to True. See b/383865404#comment3 for details.
+      # enable_hlfb=True,
   )
   return config
 

ai_edge_torch/generative/test/test_model_conversion_large.py

@@ -232,7 +232,7 @@ class TestModelConversion(googletest.TestCase):
       ai_edge_torch.config.in_oss,
       reason="tests with custom ops are not supported in oss",
   )
-  def disabled_test_paligemma1(self):
+  def test_paligemma1(self):
     self._test_paligemma_model(
         decoder.Decoder, decoder.get_fake_decoder_config, atol=1e-3, rtol=1e-5
     )
@@ -241,7 +241,7 @@ class TestModelConversion(googletest.TestCase):
       ai_edge_torch.config.in_oss,
       reason="tests with custom ops are not supported in oss",
   )
-  def disabled_test_paligemma2(self):
+  def test_paligemma2(self):
     self._test_paligemma_model(
         decoder2.Decoder2,
         decoder2.get_fake_decoder2_config,

ai_edge_torch/lowertools/common_utils.py

@@ -95,9 +95,9 @@ def _get_states(
     signatures: list[signature_module.Signature],
 ):
   for exported_program, signature in zip(exported_programs, signatures):
-    args, _ = exported_program.example_inputs
+    args, kwargs = exported_program.example_inputs
     # Calling this to get **all** the state including model buffers.
-    _flat_input_args = exported_program._graph_module_flat_inputs(args, {})
+    _flat_input_args = exported_program._graph_module_flat_inputs(args, kwargs)
     for tensor, input_spec in zip(
         _flat_input_args, exported_program.graph_signature.input_specs
     ):

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.3.0.dev20250116"
+__version__ = "0.3.0.dev20250117"

{ai_edge_torch_nightly-0.3.0.dev20250116.dist-info → ai_edge_torch_nightly-0.3.0.dev20250117.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20250116
+Version: 0.3.0.dev20250117
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI

{ai_edge_torch_nightly-0.3.0.dev20250116.dist-info → ai_edge_torch_nightly-0.3.0.dev20250117.dist-info}/RECORD

@@ -3,12 +3,12 @@ ai_edge_torch/_config.py,sha256=PKtOtBOup-cM0wBdQxby6HzuhLhIC3oq-TBG8FF4znE,2161
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/fx_pass_base.py,sha256=518ziQ0TUxqum2qZXqlD8qr65pHPh8ZNLnwFC6zvK3k,4253
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=G-QSz8n-nFkUu7vA6xPORcsrqKhus5kptuw7NkK8Au4,706
+ai_edge_torch/version.py,sha256=xfIDXOyS2Ghdmd-YYXjVjsHuMh4G95I_J1Du3sMIue4,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/conversion.py,sha256=_PoH0E1gbbsWhLGwDRwUtW2G_IgNzNF7pKQbn9ct6-4,5778
+ai_edge_torch/_convert/conversion.py,sha256=pSDY0CzZQP_jAMjSfQ1O7Ud_AF5ZDeDF-nE3nAu_hoo,5815
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
 ai_edge_torch/_convert/converter.py,sha256=075F8LRewk_033Ebsnft7FJr3KgtIbtZ_-8udIPy6ho,9980
-ai_edge_torch/_convert/signature.py,sha256=rGpBNss3Y9FCCCcdBwDo16KqavJi8N5P0M_6Wnmoz6Y,2320
+ai_edge_torch/_convert/signature.py,sha256=-YKJdLk-eNEHfhdPCtcQVtZf915SoVePEFxKXPPf16c,2572
 ai_edge_torch/_convert/to_channel_last_io.py,sha256=_31phf7TYgZY2ftpNbrdlB1RhDium1lz_BXEQ6IsMFc,2893
 ai_edge_torch/_convert/fx_passes/__init__.py,sha256=NVe-eGcm7j8jZpP2pcMhC8j5dVjgR1pPzyXhHdvKH4E,1267
 ai_edge_torch/_convert/fx_passes/build_aten_composite_pass.py,sha256=doaww8KqrgRTD5LotBVAIRFsEqzPn9R5lcGehBJOczA,9098
@@ -47,7 +47,7 @@ ai_edge_torch/generative/examples/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIX
 ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py,sha256=8HJi0cutxPstafVNs2LfBKdUzufVucje1Vrfjw_RS_g,2527
 ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=MX8fZhJJPZ5IoMiNHX0tLkRpHYqVuh4qhW0rkeIfmYw,2529
 ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=w8oWYibZzvEvCDyp39EYyAWmjgJljhzdYPyFCfAWxZA,3497
-ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=F1EG_7qPKFXtSgq5D6FYBASyvDX7N6wSQS0qdRWJzMQ,10392
+ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=CMkkTd_vO_Ej1SnmXIB0xqjRoArELOkyJ9uqjilpQeI,10298
 ai_edge_torch/generative/examples/gemma/verify_gemma1.py,sha256=ip-Gmk4CI5f0GWSdAIdrectxQWJ0t328KCsA4nfHuGg,1736
 ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=IoBhEMwH07-tFm5-U6F2hpCsI8xynglhq1x9tIOdaPQ,1322
 ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=tR8RflXocDZqvuStyw9aFlzuiTllEC8rNnjrxms6_Is,5727
@@ -65,8 +65,8 @@ ai_edge_torch/generative/examples/openelm/verify.py,sha256=VkigoqhAr8ew95neb3Tif
 ai_edge_torch/generative/examples/paligemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=scLsguzzuHfKYDWUd2uZkKYVRzdAbQHLd-kPam8QwvM,3004
 ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=g0Fbtf9WigOzQij7W1ksUca4eZTwVdCO2RcuFO2GD3M,5439
-ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=Z-vTX8dvW_UeAzCuiPOXtTqJiDDW3V1NYhOlqydpYDw,6477
-ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=yKPWG8aBp-GuzeyQntlzwTTcGBBjvUywVGRjnlNprmo,5574
+ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=craPUFxlBniBz9a0Jc7VjK01jROMg5a47xJiEA1brnw,6430
+ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=u4hEHjvLaMu-UnRrISOFXKMEJIMSLa9CfpjjmSIrlSY,5731
 ai_edge_torch/generative/examples/paligemma/paligemma.py,sha256=CEMG9gh51ev1KXPew927a6nfampiXX9bL6m-25tNYN8,6340
 ai_edge_torch/generative/examples/paligemma/verify.py,sha256=KT3Ruy40tSESxQuy-Sw01NAI3zId1BZr6Bp7FZj1wZk,5622
 ai_edge_torch/generative/examples/paligemma/verify_decoder.py,sha256=al5wMPWri4IRVWrLmCplPi6uoCzwh0vBHMGnCt-XUqo,2690
@@ -145,7 +145,7 @@ ai_edge_torch/generative/test/test_kv_cache.py,sha256=2AulHBS3hC4b_68PNNBkRVOryp
 ai_edge_torch/generative/test/test_loader.py,sha256=9mQUeeZKOVApOWSWl2cN9c10axZjMKM1-0Zd823CCS4,3449
 ai_edge_torch/generative/test/test_lora.py,sha256=6QIM6RLTc2HrodGpp_aS3OxM9Rco2KAzEnYgotkg41M,5310
 ai_edge_torch/generative/test/test_model_conversion.py,sha256=jfqkECCX7XKHeBAuDXrkwQJf0vM72eG3LMc5rluha84,6191
-ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=bBcey-aD4L_TwKRrrM81bN2VQoJjPPC84Rv4o3WOc34,12491
+ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=yzMgXkiZxHUF_xz0UR3kD3x74ELsmJetbQnmv7-9gyQ,12473
 ai_edge_torch/generative/test/test_quantize.py,sha256=bEJMhpQ9bIDUZVBXTW888728FcH-i3SyE4JSZZUgU0A,6071
 ai_edge_torch/generative/test/utils.py,sha256=tF6aCfAGJnc9dmzCnZCEOuKNVimfWOqscv9og0DDLHU,2656
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
@@ -166,7 +166,7 @@ ai_edge_torch/hlfb/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6
 ai_edge_torch/hlfb/test/test_mark_pattern.py,sha256=-5UqJyk__1YbUNGuxi4b2sn0CED0W-G337AXwxPGdEs,5567
 ai_edge_torch/lowertools/__init__.py,sha256=A8WBXvWtuFYYWtNTqPD7waVntLaSVAnSMwx5ugjZBIw,761
 ai_edge_torch/lowertools/_shim.py,sha256=Mbg16tnCVK0YsHowfbpqpNX1qySuMLvpGI_-I5SIrG0,3276
-ai_edge_torch/lowertools/common_utils.py,sha256=Z7p-ivOHtddktpnHrlDm_dSoTxJOdEjFXIGQbzjgwQo,4504
+ai_edge_torch/lowertools/common_utils.py,sha256=4HQtquPZ6oiId8vR_1ykW_uK4ELnyo5zo3MlX1QYW4c,4513
 ai_edge_torch/lowertools/odml_torch_utils.py,sha256=dxg2pBuVhSZeY2Ouc0F6nHiZilWZmpSPA7I8kGqSkVI,8282
 ai_edge_torch/lowertools/test_utils.py,sha256=mdxTlhqHABZEQ_GEmPFCL8LIAWtqRtYZUGdSY1ieZjw,1949
 ai_edge_torch/lowertools/torch_xla_utils.py,sha256=1EytIw2R6dthhLhf69wN1L9BaQTeybCD0wga-PhHcMI,9518
@@ -206,8 +206,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9
 ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.3.0.dev20250116.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20250116.dist-info/METADATA,sha256=BBIWyQ45p6eKiiajHicFpac_J0nPnyAeviNewQm5_pc,1966
-ai_edge_torch_nightly-0.3.0.dev20250116.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20250116.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20250116.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20250117.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20250117.dist-info/METADATA,sha256=5XJm1sJgKeIZBNGARZY0DOmuJB04moEM7GsVarmGwwU,1966
+ai_edge_torch_nightly-0.3.0.dev20250117.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20250117.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20250117.dist-info/RECORD,,