ai-edge-torch-nightly 0.4.0.dev20250406__py3-none-any.whl → 0.5.0.dev20250408__py3-none-any.whl

ai_edge_torch/odml_torch/export.py

@@ -209,7 +209,10 @@ class MlirLowered:
 
   def get_text(self, enable_debug_info=False):
     return str(
-        self.module.operation.get_asm(enable_debug_info=enable_debug_info)
+        self.module.operation.get_asm(
+            enable_debug_info=enable_debug_info,
+            large_elements_limit=16,
+        )
     )
 
   @property
@@ -326,8 +329,24 @@ def _convert_q_dq_per_channel_args_to_list(
 
 def exported_program_to_mlir(
     exported_program: torch.export.ExportedProgram,
+    *,
+    ir_context: ir.Context | None = None,
+    _pre_lower_pass: (
+        Callable[[torch.export.ExportedProgram], None] | None
+    ) = None,
 ) -> MlirLowered:
-  """Lower the exported program to MLIR."""
+  """Lower the exported program to MLIR.
+
+  Args:
+    exported_program: The exported program to lower.
+    ir_context: The MLIR context to use. If not provided, a new context will be
+      created.
+    _pre_lower_pass: A function to run on exported program before lowering.
+
+  Returns:
+    The lowered MLIR module, metadata, and weight tensors bundle from exported
+    program.
+  """
   exported_program = fx_infra.safe_run_decompositions(
       exported_program,
       fx_infra.decomp.pre_lower_decomp(),
@@ -340,10 +359,16 @@ def exported_program_to_mlir(
   # Do not call run_decompositions after applying the passes.
   _convert_q_dq_per_channel_args_to_list(exported_program)
 
-  with export_utils.create_ir_context() as context, ir.Location.unknown():
+  if _pre_lower_pass:
+    _pre_lower_pass(exported_program)
+
+  if not ir_context:
+    ir_context = export_utils.create_ir_context()
+
+  with ir_context, ir.Location.unknown():
 
     module = ir.Module.create()
-    lctx = LoweringContext(context, module)
+    lctx = LoweringContext(ir_context, module)
     interpreter = LoweringInterpreter(exported_program.graph_module, lctx)
     ir_flat_inputs, export_flat_args, tensor_metas = _build_flat_inputs(
         exported_program
@@ -382,7 +407,6 @@ def exported_program_to_mlir(
 
     main_func.attributes["sym_visibility"] = ir.StringAttr.get("public")
     temp_func.erase()
-
     module.operation.verify()
 
   input_signature = []
@@ -422,5 +446,5 @@ def exported_program_to_mlir(
       for tensor_meta in _get_output_metas(exported_program)
   ]
   return MlirLowered(
-      context, module, state_dict, input_signature, output_signature
+      ir_context, module, state_dict, input_signature, output_signature
   )

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.4.0.dev20250406"
+__version__ = "0.5.0.dev20250408"

{ai_edge_torch_nightly-0.4.0.dev20250406.dist-info → ai_edge_torch_nightly-0.5.0.dev20250408.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.4.0.dev20250406
+Version: 0.5.0.dev20250408
 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
@@ -25,6 +25,7 @@ License-File: LICENSE
 Requires-Dist: numpy
 Requires-Dist: scipy
 Requires-Dist: safetensors
+Requires-Dist: multipledispatchtransformerskagglehub
 Requires-Dist: tabulate
 Requires-Dist: torch>=2.4.0
 Requires-Dist: tf-nightly>=2.19.0.dev20250101

{ai_edge_torch_nightly-0.4.0.dev20250406.dist-info → ai_edge_torch_nightly-0.5.0.dev20250408.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=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=JDO-IgMGe7pl0KjfHjYEa55njUvjLd-nMF0UWJ0_6AY,706
+ai_edge_torch/version.py,sha256=Lf4c2aVfixNX2KTgdqQTLOGBdi0vVxNOkJuNt4SvQ8c,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=gpXQnifODU-mWxkUZw_3ov1lEYBw1SPVIcqj5k7pTGo,5550
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -210,7 +210,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=7l8R0DEq_vfns8iWpruMlIyaIKZAFzoAy369-7iRrl0,14164
+ai_edge_torch/odml_torch/export.py,sha256=rxsyVagQgb-DDIVtwZwSTSVFINqwIZleOOfmPkBoPKg,14817
 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
@@ -243,8 +243,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.4.0.dev20250406.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.4.0.dev20250406.dist-info/METADATA,sha256=ePIGy44ydD_RrrDoZipgy4rPG-K-xCkusBrXRBlVg3U,1966
-ai_edge_torch_nightly-0.4.0.dev20250406.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.4.0.dev20250406.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.4.0.dev20250406.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250408.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250408.dist-info/METADATA,sha256=-Bw-LUn9l-B66aMZiFiUiYBifr1B6Fr86LU8KXtBieo,2019
+ai_edge_torch_nightly-0.5.0.dev20250408.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250408.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250408.dist-info/RECORD,,