ai-edge-torch-nightly 0.3.0.dev20241106__py3-none-any.whl → 0.3.0.dev20241108__py3-none-any.whl

ai_edge_torch/fx_pass_base.py

@@ -76,7 +76,7 @@ def run_passes(
             range_constraints=exported_program.range_constraints,
             module_call_graph=exported_program.module_call_graph,
             example_inputs=exported_program.example_inputs,
-            verifier=exported_program.verifier,
+            verifiers=exported_program.verifiers,
             constants=exported_program.constants,
         )
   return exported_program

ai_edge_torch/odml_torch/_torch_future.py

@@ -59,3 +59,17 @@ def graph_module_flat_inputs(ep: torch.export.ExportedProgram, args, kwargs):
     ordered_tensor_constants = tuple()
 
   return (*param_buffer_values, *flat_args, *ordered_tensor_constants)
+
+
+# TODO(b/331481564): Replace this with CanonicalizePass + run_decomposition
+def safe_run_decompositions(exported_program, decomp_table=None):
+  for node in exported_program.graph.nodes:
+    if node.target == torch.ops.aten.view.default:
+      # Passes or torch.export may generate aten.view nodes not respecting the
+      # tensor memory format. Changes all the aten.view to torch.reshape
+      # for retracing. If the input memory format is already contiguous,
+      # retracing in run_decomposition below would decompose torch.reshape
+      # back to one aten.view.
+      node.target = lambda self, size: torch.reshape(self.contiguous(), size)
+
+  return exported_program.run_decompositions(decomp_table)

ai_edge_torch/odml_torch/export.py

@@ -262,13 +262,13 @@ def exported_program_to_mlir(
     exported_program: torch.export.ExportedProgram,
 ) -> MlirLowered:
   """Lower the exported program to MLIR."""
-  exported_program = exported_program.run_decompositions(
-      lowerings.decompositions()
+  exported_program = _torch_future.safe_run_decompositions(
+      exported_program, lowerings.decompositions()
   )
 
   _convert_i64_to_i32(exported_program)
-  exported_program = exported_program.run_decompositions(
-      lowerings.decompositions()
+  exported_program = _torch_future.safe_run_decompositions(
+      exported_program, lowerings.decompositions()
   )
 
   with export_utils.create_ir_context() as context, ir.Location.unknown():

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241106.dist-info → ai_edge_torch_nightly-0.3.0.dev20241108.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20241106
+Version: 0.3.0.dev20241108
 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
@@ -29,7 +29,7 @@ Requires-Dist: safetensors
 Requires-Dist: tabulate
 Requires-Dist: torch>=2.4.0
 Requires-Dist: torch-xla>=2.4.0
-Requires-Dist: tf-nightly==2.19.0.dev20241001
+Requires-Dist: tf-nightly>=2.19.0.dev20241001
 Requires-Dist: ai-edge-litert-nightly
 Requires-Dist: ai-edge-quantizer-nightly
 

{ai_edge_torch_nightly-0.3.0.dev20241106.dist-info → ai_edge_torch_nightly-0.3.0.dev20241108.dist-info}/RECORD

@@ -1,9 +1,9 @@
 ai_edge_torch/__init__.py,sha256=48qP37uHT90YPs4eIUQxCiWVwqGEX3idCUs6mQKvX1U,1168
 ai_edge_torch/config.py,sha256=FMWeCH2b7HYILBvaI1iZNnYCO4WAhDOwBZBmIE-xrF0,909
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
-ai_edge_torch/fx_pass_base.py,sha256=SrYveglaiA_DXPoRBqSXClWM1q7853I5ujRorq_MV0M,4251
+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=sQ3ZJv6ITWXDYt620D2xXSV4hpsToWri5zvksNUd0SI,706
+ai_edge_torch/version.py,sha256=sBOl0mgVPJtokiP8qTbTtY0R_qIaF0KNiALh7P3AJEk,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=HwzfRx_DX5TLtPqwEH1_NOm38_INvHzHl4_mX67KOdQ,5448
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -155,9 +155,9 @@ ai_edge_torch/lowertools/test_utils.py,sha256=bPgc2iXX16KYtMNvmsRdKfrCY6UJmcfitf
 ai_edge_torch/lowertools/torch_xla_utils.py,sha256=XGZE0vZG9WSQT-6dFmPlU8W89z8rfXPRGjuZeuhXCIw,9205
 ai_edge_torch/lowertools/translate_recipe.py,sha256=ymkBpFqAUiupRWqrPOWiVphKcXR1K5vHK0RjgBFtxlE,5652
 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_future.py,sha256=AJ0klpsbu2ZBTfiZlqSOoaYzBVITt40a1fYN8xKkEPw,3044
 ai_edge_torch/odml_torch/_torch_library.py,sha256=Lw1gqL2HWNRspdTwNhIkYAHDyafHedHtkXyKKxn-Wss,805
-ai_edge_torch/odml_torch/export.py,sha256=q6cvB3RAKL7hlQFNGqsE3u-NqWWPSzj-8M38u8loSpk,11544
+ai_edge_torch/odml_torch/export.py,sha256=4xwrsDeOAgzoB9m7EeNsBj6dC5Ajtn5aKDRQkdHxa-o,11584
 ai_edge_torch/odml_torch/export_utils.py,sha256=q84U69ZQ82hLXw-xncJ8IW-K71Xux-NWlzZTs7hdZWA,5127
 ai_edge_torch/odml_torch/tf_integration.py,sha256=lTFJPPEijLPFmn6qq2jbpVTQOo0YaOTK36kK6rCiyIE,5956
 ai_edge_torch/odml_torch/composite/__init__.py,sha256=71GM_gDZxJyo38ZSoYSwhZX3xKA9rknO93JS9kw9w_c,778
@@ -186,8 +186,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.dev20241106.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241106.dist-info/METADATA,sha256=qP9cLptTjaybmmaN6zeEgnSJ2cvIzt2L_B_4X5wXYXs,1897
-ai_edge_torch_nightly-0.3.0.dev20241106.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20241106.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241106.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241108.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241108.dist-info/METADATA,sha256=gp2VN_X4YPdK8axZYIhqafgiJhCwfiN_tOWT-yL3lW0,1897
+ai_edge_torch_nightly-0.3.0.dev20241108.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+ai_edge_torch_nightly-0.3.0.dev20241108.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241108.dist-info/RECORD,,