ai-edge-torch-nightly 0.5.0.dev20250417__py3-none-any.whl → 0.5.0.dev20250419__py3-none-any.whl

ai_edge_torch/_convert/conversion.py

@@ -35,14 +35,12 @@ def _run_convert_passes(
   )
 
   passes = [
+      fx_passes.CastInputsBf16ToF32Pass(),
       fx_passes.BuildInterpolateCompositePass(),
-      fx_passes.CanonicalizePass(),
       fx_passes.OptimizeLayoutTransposesPass(),
       fx_passes.CanonicalizePass(),
       fx_passes.BuildAtenCompositePass(),
       fx_passes.RemoveNonUserOutputsPass(),
-      fx_passes.CastInputsBf16ToF32Pass(),
-      fx_passes.CanonicalizePass(),
   ]
 
   # Debuginfo is not injected automatically by odml_torch. Only inject

ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py

@@ -48,6 +48,7 @@ def main(_):
       pixel_values_size=torch.Size(
           [1, config.channels, config.image_size, config.image_size]
       ),
+      pixel_seq_len=(config.image_size // config.patch_size) ** 2,
       quantize=flags.FLAGS.quantize,
       config=pytorch_model.config.decoder_config,
       export_config=ExportConfig(),

ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py

@@ -43,6 +43,9 @@ def main(_):
   )
 
   grid_thw = pytorch_model.image_encoder.get_grid_thw()
+  spatial_merge_size = (
+      pytorch_model.config.image_encoder_config.spatial_merge_size
+  )
   converter.convert_to_tflite(
       pytorch_model,
       output_path=flags.FLAGS.output_path,
@@ -51,6 +54,10 @@ def main(_):
       pixel_values_size=(
           pytorch_model.image_encoder.get_pixel_values_size(grid_thw)
       ),
+      pixel_seq_len=(
+          (grid_thw[0][1] // spatial_merge_size)
+          * (grid_thw[0][2] // spatial_merge_size)
+      ),
       quantize=flags.FLAGS.quantize,
       config=pytorch_model.config.decoder_config,
       export_config=ExportConfig(),

ai_edge_torch/generative/fx_passes/__init__.py

@@ -23,8 +23,5 @@ def run_generative_passes(
 ) -> torch.export.ExportedProgram:
   return fx_infra.run_passes(
       exported_program,
-      [
-          RemoveSDPACompositeZeroMaskPass(),
-          CanonicalizePass(),
-      ],
+      [RemoveSDPACompositeZeroMaskPass()],
   )

ai_edge_torch/generative/utilities/converter.py

@@ -57,7 +57,7 @@ def define_conversion_flags(model_name: str):
   )
   flags.DEFINE_string(
       'output_name_prefix',
-      f'{model_name}',
+      model_name,
       'The prefix of the output tflite model name.',
   )
   flags.DEFINE_multi_integer(
@@ -91,6 +91,7 @@ def convert_to_tflite(
     output_name_prefix: str,
     prefill_seq_len: Union[int, list[int]],
     pixel_values_size: torch.Size = None,
+    pixel_seq_len: int = 0,
     quantize: bool = True,
     config: cfg.ModelConfig = None,
     lora_ranks: Optional[list[int]] = None,
@@ -133,12 +134,18 @@ def convert_to_tflite(
         use. If a list, the model will have multiple prefill signatures.
       pixel_values_size (torch.Size, optional): The size of pixel values to pass
         to the model. If None, the model is not expected to take pixel values.
+      pixel_seq_len (int, optional): The length of pixel tokens, or pixel
+        embeddings generated by the image encoder with pixel values. The actual
+        length of prefill_seq_len will be added by pixel_seq_len when pixel
+        values are passed.
       quantize (bool, optional): Whether the model should be quanized. Defaults
         to True.
       config (cfg.ModelConfig, optional): The model config used to configure KV
         cache. If None, it uses the config of the pytorch_model.
       lora_ranks (list[int], optional): The ranks of the LORA layers. If None,
         no LoRA signatures will be added.
+      export_config (ExportConfig, optional): The export configuration. If None,
+        it uses the default export configuration.
   """
   # pylint: disable=protected-access
   torch._dynamo.config.cache_size_limit = 64
@@ -173,6 +180,7 @@ def convert_to_tflite(
       output_file,
       prefill_seq_lens,
       pixel_values_size,
+      pixel_seq_len,
       quantize,
       config,
       loras,
@@ -185,6 +193,7 @@ def _export_helper(
     output_file: str,
     prefill_seq_lens: list[int],
     pixel_values_size: torch.Size,
+    pixel_seq_len: int,
     quantize: bool,
     config: cfg.ModelConfig,
     loras: list[None | lora_utils.LoRA],
@@ -197,11 +206,18 @@ def _export_helper(
     prefill_tokens_list.append(torch.full((1, seq_len), 0, dtype=torch.int))
     prefill_input_pos_list.append(torch.arange(0, seq_len, dtype=torch.int))
 
-  prefill_pixel_values = (
-      torch.full(pixel_values_size, 0, dtype=torch.float32)
-      if pixel_values_size
-      else None
-  )
+  prefill_pixel_values = None
+  prefill_tokens_list_with_pixel = []
+  prefill_input_pos_list_with_pixel = []
+  if pixel_values_size is not None:
+    prefill_pixel_values = torch.full(pixel_values_size, 0, dtype=torch.float32)
+    for seq_len in prefill_seq_lens:
+      prefill_tokens_list_with_pixel.append(
+          torch.full((1, seq_len + pixel_seq_len), 0, dtype=torch.int)
+      )
+      prefill_input_pos_list_with_pixel.append(
+          torch.arange(0, seq_len + pixel_seq_len, dtype=torch.int)
+      )
 
   if export_config.prefill_mask is None:
     prefill_masks = None
@@ -238,13 +254,11 @@ def _export_helper(
   for lora in loras:
     for i in range(len(prefill_seq_lens)):
       prefill_seq_len = prefill_seq_lens[i]
-      prefill_tokens = prefill_tokens_list[i]
-      prefill_input_pos = prefill_input_pos_list[i]
       prefill_signature_name = f'prefill_{prefill_seq_len}'
 
       sample_kwargs = {
-          'tokens': prefill_tokens,
-          'input_pos': prefill_input_pos,
+          'tokens': prefill_tokens_list[i],
+          'input_pos': prefill_input_pos_list[i],
           'kv_cache': prefill_kv,
       }
       if prefill_masks is not None:
@@ -261,13 +275,13 @@ def _export_helper(
       )
 
       if prefill_pixel_values is not None:
+        sample_kwargs['tokens'] = prefill_tokens_list_with_pixel[i]
+        sample_kwargs['input_pos'] = prefill_input_pos_list_with_pixel[i]
+        sample_kwargs['pixel_values'] = prefill_pixel_values
         converter.add_signature(
             prefill_signature_name + '_pixel',
             mod,
-            sample_kwargs={
-                **sample_kwargs,
-                'pixel_values': prefill_pixel_values,
-            },
+            sample_kwargs=sample_kwargs,
         )
 
     sample_kwargs = {

ai_edge_torch/odml_torch/export.py

@@ -264,6 +264,8 @@ def _convert_i64_to_i32(exported_program: torch.export.ExportedProgram):
     exported_program: The exported program to apply the pass.
   """
 
+  is_modified = False
+
   def in_i32(x: int):
     return -2147483648 <= x <= 2147483647
 
@@ -271,6 +273,7 @@ def _convert_i64_to_i32(exported_program: torch.export.ExportedProgram):
     return torch.ops.aten._to_copy.default(x, dtype=torch.int32)
 
   def rewrite_arange(node: torch.fx.Node):
+    nonlocal is_modified
     tensor_meta = node.meta.get("tensor_meta", None)
     if not tensor_meta:
       return
@@ -282,12 +285,14 @@ def _convert_i64_to_i32(exported_program: torch.export.ExportedProgram):
       return
     op = node.target
     node.target = lambda *args, **kwargs: to_int32(op(*args, **kwargs))
+    is_modified = True
 
   graph_module = exported_program.graph_module
   for node in graph_module.graph.nodes:
 
     if node.target == torch.ops.aten.arange.start_step:
       rewrite_arange(node)
+  return is_modified
 
 
 # TODO(b/331481564) Make this a ai_edge_torch FX pass.
@@ -351,9 +356,9 @@ def exported_program_to_mlir(
       exported_program,
       fx_infra.decomp.pre_lower_decomp(),
   )
-  _convert_i64_to_i32(exported_program)
-  # Run decompositions for retracing and cananicalization.
-  exported_program = fx_infra.safe_run_decompositions(exported_program, {})
+  if _convert_i64_to_i32(exported_program):
+    # Run decompositions for retracing and cananicalization, if modified.
+    exported_program = fx_infra.safe_run_decompositions(exported_program, {})
 
   # Passes below mutate the exported program to a state not executable by torch.
   # Do not call run_decompositions after applying the passes.

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.5.0.dev20250417.dist-info → ai_edge_torch_nightly-0.5.0.dev20250419.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250417
+Version: 0.5.0.dev20250419
 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.5.0.dev20250417.dist-info → ai_edge_torch_nightly-0.5.0.dev20250419.dist-info}/RECORD

@@ -2,9 +2,9 @@ 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=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=kwb6M7GEr85K6sLrsbI9sNCggXojl_5TX9GeVCyP9OI,706
+ai_edge_torch/version.py,sha256=SG1Sn0KkGdZyTIYeY_Rw8sRC6xtCmFSkF15xymi-Eho,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/conversion.py,sha256=GPDsXhfECjDzOut4vh_d9qWcyfpxobFMBTsC7MyJbM0,5557
+ai_edge_torch/_convert/conversion.py,sha256=0gpwEjlTue5RttDerzM5SVOUnY8g16444yL2YIFBx-E,5485
 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=-YKJdLk-eNEHfhdPCtcQVtZf915SoVePEFxKXPPf16c,2572
@@ -84,7 +84,7 @@ ai_edge_torch/generative/examples/openelm/convert_to_tflite.py,sha256=Hgp31zIQdJ
 ai_edge_torch/generative/examples/openelm/openelm.py,sha256=sIJ8Ie1oxFrJM-1jvv2ukiJbQOTIUGuMEZvmwZbt3n0,4556
 ai_edge_torch/generative/examples/openelm/verify.py,sha256=4W26ZtPF5Cb9mpHYuRM4b2QB_4W76zf4WV36KzexVjs,2446
 ai_edge_torch/generative/examples/paligemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=2GLE4empjc8IANssR02ECFUqhdUNJV_OVHCf1UXKL8Y,1956
+ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=7HHXkC-IIu7ieBvBI4RlXs_oITz7R8a6YVYQskAs_Uk,2023
 ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=G1dwtWp_v77AI3uyIY-8g6qRP2tRH3CIKjJTeYNqFPU,5511
 ai_edge_torch/generative/examples/paligemma/decoder2.py,sha256=Z-SKdb0dd8uWT1d-FRwFx5-tJEqpdrQwiIZnFRhOtVo,6060
 ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=SvuR97sjkBtfkerH7Hu1UXB8kCFLpEATNbPfCbNAyfo,5614
@@ -108,7 +108,7 @@ ai_edge_torch/generative/examples/qwen/convert_to_tflite.py,sha256=om3lXL1RnA87P
 ai_edge_torch/generative/examples/qwen/qwen.py,sha256=Zi_qiQ1JPokXZ95jgSEnQp3F-LKzFCvWvFLKhJjnASo,4199
 ai_edge_torch/generative/examples/qwen/verify.py,sha256=9_AyEJTeUfvhhID64Rto2bflFPyXMFokdQLsseLUMiI,2775
 ai_edge_torch/generative/examples/qwen_vl/__init__.py,sha256=JaAnrFoXTl3RJX97XspklkTyqOHVyAgRJsZtzNDd10c,671
-ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py,sha256=OcE2-8lqAukoK5hM1sqdgfXU37kxWQ84racweNAdjyk,1995
+ai_edge_torch/generative/examples/qwen_vl/convert_to_tflite.py,sha256=yVebRatt2SLCsGvrYTBXOM-0S2REhkpikHTyy5MCjUw,2222
 ai_edge_torch/generative/examples/qwen_vl/decoder.py,sha256=7RFM25tDj_b0FkpSv8RUWir8K8v9p2jMtwZmP4VAUhw,4474
 ai_edge_torch/generative/examples/qwen_vl/image_encoder.py,sha256=nHzBe_YSPnUe1d5i09v4bePQomVifzJNeUjRfprmxC0,14878
 ai_edge_torch/generative/examples/qwen_vl/qwen_vl.py,sha256=mfLFrT8NPEPh9CqlJYHwh-I2y6ST7hH_vEmbZYartHQ,7764
@@ -147,7 +147,7 @@ ai_edge_torch/generative/examples/tiny_llama/__init__.py,sha256=hHLluseD2R0Hh4W6
 ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=LPxg7mAJ_aAUIx6eE5bxixPA8Ep9Vul0CWJoNcrD5oE,1565
 ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=mhJ18rb9sxrYRzv1YSzhbNs97oUZck99avZDcUO2oV8,2800
 ai_edge_torch/generative/examples/tiny_llama/verify.py,sha256=LRu6PSw7Lqu6HGbv1tO2i0nUCqe-VkRgboA10VZ7KNg,2431
-ai_edge_torch/generative/fx_passes/__init__.py,sha256=4rFrppMRKlTwwZeX1ON_cdp4yUqoTOES161IZQkJF6c,1143
+ai_edge_torch/generative/fx_passes/__init__.py,sha256=PFSMsA1vfBfrV9ssBCkYJNl8Hx_bLdWjN01iyjPM5jE,1094
 ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=myGjal5A8yIBoqgArd2k40rZmCgD1Ya369KR7182bhI,2129
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/attention.py,sha256=wLZ1jgUlcODBWgK3hnnhclHuuQDqYuGOZdYAI9EooOM,13247
@@ -186,7 +186,7 @@ ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=-v2Vj7Qdd3Gy
 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
-ai_edge_torch/generative/utilities/converter.py,sha256=swtz69oyMOxSaCEYST_Gzd5sjGZ1qOBAfd_0xl207Nk,9766
+ai_edge_torch/generative/utilities/converter.py,sha256=LtBHjnslhL-uf4sDRoC8JIbbUD73g0QW3FiWsHUdV1g,10631
 ai_edge_torch/generative/utilities/export_config.py,sha256=8-795nyd3M34LkGhgW7hwHlJyTc2Oz1iipHK8yBhdFs,1633
 ai_edge_torch/generative/utilities/loader.py,sha256=7p__m2JryWphGlYOuRxdoT4id4_tWJEVOV7y2X4H-Ak,13737
 ai_edge_torch/generative/utilities/model_builder.py,sha256=ZYX1TxpFdj573du2QCyHJlFjx4q1m12R74fp4Gwl92A,6343
@@ -212,7 +212,7 @@ ai_edge_torch/lowertools/translate_recipe.py,sha256=ymkBpFqAUiupRWqrPOWiVphKcXR1
 ai_edge_torch/odml_torch/__init__.py,sha256=S8jOzE9nLof-6es3XDiGJRN-9H_XTxsVm9dE7lD3RWo,812
 ai_edge_torch/odml_torch/_torch_future.py,sha256=jSYHf1CMTJzMizPMbu2b39hAt0ZTR6gQLq67GMe9KTo,2336
 ai_edge_torch/odml_torch/_torch_library.py,sha256=Lw1gqL2HWNRspdTwNhIkYAHDyafHedHtkXyKKxn-Wss,805
-ai_edge_torch/odml_torch/export.py,sha256=rxsyVagQgb-DDIVtwZwSTSVFINqwIZleOOfmPkBoPKg,14817
+ai_edge_torch/odml_torch/export.py,sha256=lbLpdGa8MDE8oWNA7aSV3tOCQ9P9I2Ox95dSPEssn-g,14930
 ai_edge_torch/odml_torch/export_utils.py,sha256=QeA37Irlty6AiIBuqmHmJgn3lqahBQ5xsh6IKRoKm1g,4774
 ai_edge_torch/odml_torch/tf_integration.py,sha256=NN29WeXmHZ0S1RPDFHUnBi2DEjMvAtwczStPYIsQ1w8,4849
 ai_edge_torch/odml_torch/composite/__init__.py,sha256=71GM_gDZxJyo38ZSoYSwhZX3xKA9rknO93JS9kw9w_c,778
@@ -245,8 +245,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=dguMa-aEi-WDPnmGBUs2IPdEmt2IVmHOELH19uiJ1uU,3014
 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.dev20250417.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250417.dist-info/METADATA,sha256=ovMriaKRgveZtN2i-cTOM2_8BuNvgf-SYNITAte1wjs,2051
-ai_edge_torch_nightly-0.5.0.dev20250417.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250417.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250417.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250419.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250419.dist-info/METADATA,sha256=FqixHlt1f3QPZdgxBHMWfF_GAD2GwOXkFvyVMP8IjpI,2051
+ai_edge_torch_nightly-0.5.0.dev20250419.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250419.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250419.dist-info/RECORD,,