ai-edge-torch-nightly 0.5.0.dev20250508__py3-none-any.whl → 0.5.0.dev20250509__py3-none-any.whl

ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py

@@ -20,7 +20,9 @@ from ai_edge_torch.generative.examples.gemma import gemma2
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
 
-flags = converter.define_conversion_flags("gemma2-2b")
+flags = converter.define_conversion_flags(
+    "gemma2-2b", default_mask_as_input=True, default_transpose_kv_cache=True
+)
 
 
 def main(_):

ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py

@@ -20,7 +20,9 @@ from ai_edge_torch.generative.examples.gemma3 import gemma3
 from ai_edge_torch.generative.utilities import converter
 from ai_edge_torch.generative.utilities import export_config
 
-flags = converter.define_conversion_flags('gemma3-1b')
+flags = converter.define_conversion_flags(
+    'gemma3-1b', default_mask_as_input=True, default_transpose_kv_cache=True
+)
 
 _MODEL_SIZE = flags.DEFINE_string(
     'model_size',

ai_edge_torch/generative/examples/gemma3/decoder.py

@@ -261,7 +261,6 @@ class Decoder(nn.Module):
       pixel_mask = self.build_pixel_mask(image_indices)
     # RoPE parameters are the same for all blocks. Use the first layer.
     attn_config = self.config.block_config(0).attn_config
-    n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
     # Different rotary base for global and local attention
     # based on attention pattern
     rope = [
@@ -305,7 +304,7 @@ class Decoder(nn.Module):
     if pixel_mask is None:
       mask = [
           self.get_local_global_attention_mask(
-              mask,
+              mask[i] if isinstance(mask, list) else mask,
               self.config.block_config(i).attn_config.attn_type,
               input_pos,
               self.config.block_config(i).attn_config.sliding_window_size,
@@ -316,7 +315,7 @@ class Decoder(nn.Module):
       pixel_mask = pixel_mask.index_select(2, input_pos)
       mask = [
           self.compose_mask(
-              mask[i],
+              mask[i] if isinstance(mask, list) else mask,
               pixel_mask,
               self.config.block_config(i).attn_config.attn_type,
           )

ai_edge_torch/generative/utilities/converter.py

@@ -42,7 +42,11 @@ class ExportableModule(torch.nn.Module):
     return self.module(*export_args, **full_kwargs)
 
 
-def define_conversion_flags(model_name: str):
+def define_conversion_flags(
+    model_name: str,
+    default_mask_as_input: bool = False,
+    default_transpose_kv_cache: bool = False,
+):
   """Defines common flags used for model conversion."""
 
   flags.DEFINE_string(
@@ -83,13 +87,13 @@ def define_conversion_flags(model_name: str):
   )
   flags.DEFINE_bool(
       'mask_as_input',
-      False,
+      default_mask_as_input,
       'If true, the mask will be passed in as input. Otherwise, mask will be '
       'built by the model internally.',
   )
   flags.DEFINE_bool(
       'transpose_kv_cache',
-      False,
+      default_transpose_kv_cache,
       'If true, the model will be converted with transposed KV cache.',
   )
   return flags

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.5.0.dev20250508"
+__version__ = "0.5.0.dev20250509"

{ai_edge_torch_nightly-0.5.0.dev20250508.dist-info → ai_edge_torch_nightly-0.5.0.dev20250509.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250508
+Version: 0.5.0.dev20250509
 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
@@ -22,6 +22,7 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: absl-py
 Requires-Dist: numpy
 Requires-Dist: scipy
 Requires-Dist: safetensors

{ai_edge_torch_nightly-0.5.0.dev20250508.dist-info → ai_edge_torch_nightly-0.5.0.dev20250509.dist-info}/RECORD

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=8sPR_5uXJA4NEE0nIwNdSl-ADOJEoR8hAgYvBQDY70Y,120
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=wxjSFq_rBSxSqbUE8E8EJTCkgvgaRLjq_ZuAM-IZpCU,5606
-ai_edge_torch/version.py,sha256=7lrbHHeWyBpqJdwFYYooOGJss4Rvg3UAdFSo9K0uzek,706
+ai_edge_torch/version.py,sha256=bmd7zA2ryjEwRkVoQtjuMQTzdZ0gresufnMFls2eWNo,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=QVugYVfbyaeBgSKKbhFzHG5oXA7t3M-40JcpcdSu6W8,5436
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -61,15 +61,15 @@ ai_edge_torch/generative/examples/deepseek/deepseek.py,sha256=yhS_i2kR0GJWpWciCt
 ai_edge_torch/generative/examples/deepseek/verify.py,sha256=iYldze-pvZGvPkkqr6zA7EmitPnH9sXkzjNVx353IcE,2403
 ai_edge_torch/generative/examples/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py,sha256=RRilUl2Ui08R9gy1Ua0jnaXNCrIJJb-oztgP62G3mX4,1526
-ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=7IlF-4NEfZAzIfkOUHR-HeCSLSUGEu7wnO52UtERCa4,1527
+ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=9ozSw2-xuf5Wfh1HeLDTP3wJxxUZmrD3An1njJPMpdI,1594
 ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=6ImjTzJcq6JoKz2Z-z8pjv5BsRu5nUeEsTK3IPs3xgI,3521
 ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=JQLLiHNVBM9jOrZqUF0EmgAwtDD0yTRlmIbLaWM7qTg,11557
 ai_edge_torch/generative/examples/gemma/verify_gemma1.py,sha256=ip-Gmk4CI5f0GWSdAIdrectxQWJ0t328KCsA4nfHuGg,1736
 ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=jhiyinOqPt5ZZjEadDRZt_wY5fiLSCpMo54PcxFaL_Q,1789
 ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=n7f2nF6Lin_tDvPs0JVldsuaBzo7pAwi5YAHAhlIxQg,6139
 ai_edge_torch/generative/examples/gemma3/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
-ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=MjkQDVynaw9C5z9ODzKfb85xW5JfxHUWBJ_Aco05FHo,1760
-ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=xGxeNKQvgyrENmUQMu0uKymL3qthvbdoxdMbAzwiLz0,15725
+ai_edge_torch/generative/examples/gemma3/convert_gemma3_to_tflite.py,sha256=JLXXn2mFEBs4DlHH_O6hpEG9KInJqsCdWy3DrgUjT1c,1827
+ai_edge_torch/generative/examples/gemma3/decoder.py,sha256=soxNVtN2fns9pQWw55ZND7dJ9RQLkFBtAteDSwZO9oY,15729
 ai_edge_torch/generative/examples/gemma3/gemma3.py,sha256=GACDBI_MsFowR8A3wAWrpzradPYe-AUgB9ZjXaVBG-s,6485
 ai_edge_torch/generative/examples/gemma3/image_encoder.py,sha256=uRoLoBWzFtQz5wFZfPCxbkvZsgPAqSkUUsV3977GbYc,5184
 ai_edge_torch/generative/examples/gemma3/verify_gemma3.py,sha256=v8oNXFICmVOtQxfO7IhZ8GnbvotEkDi9lzYHjoQyOso,2464
@@ -192,7 +192,7 @@ ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=vQWmpzMkJ2hP
 ai_edge_torch/generative/test/test_quantize.py,sha256=TG6vTF9yOZWe2wW7v8-hmuaQoODwJC1Z-2d5xv3zgfI,7389
 ai_edge_torch/generative/test/utils.py,sha256=tF6aCfAGJnc9dmzCnZCEOuKNVimfWOqscv9og0DDLHU,2656
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
-ai_edge_torch/generative/utilities/converter.py,sha256=d0JOWN5l2vbvt8RzFFiRoulkWiejyEZ21xKv5LdLIyc,11675
+ai_edge_torch/generative/utilities/converter.py,sha256=u-FViLhHrbcO-GYfcGXTY28Kf-N682j-sao7LkdZbJ0,11806
 ai_edge_torch/generative/utilities/export_config.py,sha256=5IvR3grlMd4mWO5c_Y4x9Fk1b1xa57MzlYNE8XUaN28,2049
 ai_edge_torch/generative/utilities/loader.py,sha256=7p__m2JryWphGlYOuRxdoT4id4_tWJEVOV7y2X4H-Ak,13737
 ai_edge_torch/generative/utilities/model_builder.py,sha256=ZYX1TxpFdj573du2QCyHJlFjx4q1m12R74fp4Gwl92A,6343
@@ -251,8 +251,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
 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.5.0.dev20250508.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250508.dist-info/METADATA,sha256=GGDJl2Fya8gLr9RIfSLCmm1K1xA3qzBrrEOy1hwR2dQ,2051
-ai_edge_torch_nightly-0.5.0.dev20250508.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250508.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250508.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250509.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250509.dist-info/METADATA,sha256=IrEO6k64_TTmFzYdwfbTfwPaNX-Lhcg3pwAV8sGw4Gk,2074
+ai_edge_torch_nightly-0.5.0.dev20250509.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250509.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250509.dist-info/RECORD,,