ai-edge-torch-nightly 0.5.0.dev20250418__py3-none-any.whl → 0.5.0.dev20250420__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/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/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.dev20250418"
+__version__ = "0.5.0.dev20250420"

{ai_edge_torch_nightly-0.5.0.dev20250418.dist-info → ai_edge_torch_nightly-0.5.0.dev20250420.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250418
+Version: 0.5.0.dev20250420
 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.dev20250418.dist-info → ai_edge_torch_nightly-0.5.0.dev20250420.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=RijgVzFc62dFf_A0Pr5jwDo3wtofAoxUcfBieTnHvIw,706
+ai_edge_torch/version.py,sha256=yXQZIt1y2LjGAGZs_dNNktRUPPQAI4FYJV4wZ1gLZWI,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
@@ -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
@@ -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.dev20250418.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250418.dist-info/METADATA,sha256=cmJPI-Zen5YAjRWQPoVwvcAwVZJGz3Jz0OuNAHs5498,2051
-ai_edge_torch_nightly-0.5.0.dev20250418.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250418.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250418.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250420.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250420.dist-info/METADATA,sha256=dqTzx2gTDf6SkPglyObeZL_KZAlGzDITMcbP3qHKINk,2051
+ai_edge_torch_nightly-0.5.0.dev20250420.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250420.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250420.dist-info/RECORD,,