ai-edge-torch-nightly 0.3.0.dev20241010__py3-none-any.whl → 0.3.0.dev20241011__py3-none-any.whl

ai_edge_torch/_convert/conversion.py

@@ -15,7 +15,7 @@
 
 import logging
 import os
-from typing import Any, Optional
+from typing import Any, Literal, Optional, Union
 
 from ai_edge_torch import fx_pass_base
 from ai_edge_torch import lowertools
@@ -73,6 +73,7 @@ def _warn_training_modules(signatures: list[signature.Signature]):
 def convert_signatures(
     signatures: list[signature.Signature],
     *,
+    strict_export: Union[Literal["auto"], bool] = True,
     quant_config: Optional[qcfg.QuantConfig] = None,
     _tfl_converter_flags: Optional[dict[str, Any]],
 ) -> model.TfLiteModel:
@@ -81,6 +82,11 @@ def convert_signatures(
   Args:
       signatures: The list of 'signature.Signature' objects containing PyTorch
         modules to be converted.
+      strict_export: Experimental `strict` arg for torch.export.export. When
+        enabled, the export function will trace the program through TorchDynamo
+        and ensure the soundness of the exported graph. When
+        strict_export="auto", the function will try to export module in both
+        modes and use the first one succeeds for downstream conversion.
       quant_config: User-defined quantization method and scheme of the model.
       _tfl_converter_flags: A nested dictionary allowing setting flags for the
         underlying tflite converter.
@@ -93,10 +99,24 @@ def convert_signatures(
 
   _warn_training_modules(signatures)
 
-  exported_programs: torch.export.torch.export.ExportedProgram = [
-      torch.export.export(
-          sig.module, sig.flat_args, dynamic_shapes=sig.dynamic_shapes
-      )
+  def export(*args, **kwargs):
+    nonlocal strict_export
+    if strict_export == "auto":
+      try:
+        return torch.export.export(*args, **kwargs, strict=True)
+      except Exception:
+        logging.warning(
+            "torch.export.export(..., strict=True) failed. Retrying with"
+            " strict=False"
+        )
+        return torch.export.export(*args, **kwargs, strict=False)
+    elif not strict_export:
+      return torch.export.export(*args, **kwargs, strict=False)
+    else:
+      return torch.export.export(*args, **kwargs, strict=True)
+
+  exported_programs: torch.export.ExportedProgram = [
+      export(sig.module, sig.flat_args, dynamic_shapes=sig.dynamic_shapes)
       for sig in signatures
   ]
 

ai_edge_torch/_convert/converter.py

@@ -15,7 +15,7 @@
 
 from __future__ import annotations
 
-from typing import Any, Optional, Tuple, Union
+from typing import Any, Literal, Optional, Tuple, Union
 
 from ai_edge_torch import model
 from ai_edge_torch._convert import conversion
@@ -102,6 +102,7 @@ class Converter:
       sample_args=None,
       sample_kwargs=None,
       *,
+      strict_export: Union[Literal["auto"], bool] = True,
       quant_config: Optional[qcfg.QuantConfig] = None,
       dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
       _ai_edge_converter_flags: Optional[dict[str, Any]] = None,
@@ -123,6 +124,11 @@ class Converter:
         with prior to conversion.
       sample_kwargs: Dict of str to tensor by which the torch module will be
         traced with prior to conversion.
+      strict_export: Experimental `strict` arg for torch.export.export. When
+        enabled, the export function will trace the program through TorchDynamo
+        and ensure the soundness of the exported graph. When
+        strict_export="auto", the function will try to export module in both
+        modes and use the first one succeeds for downstream conversion.
       quant_config: User-defined quantization method and scheme of the model.
       dynamic_shapes: Optional dict or tuple that specify dynamic shape
         specifications for each input in original order. See
@@ -162,6 +168,7 @@ class Converter:
         )
     return conversion.convert_signatures(
         self._signatures,
+        strict_export=strict_export,
         quant_config=quant_config,
         _tfl_converter_flags=_ai_edge_converter_flags,
     )
@@ -205,6 +212,7 @@ def convert(
     sample_args=None,
     sample_kwargs=None,
     *,
+    strict_export: Union[Literal["auto"], bool] = True,
     quant_config: Optional[qcfg.QuantConfig] = None,
     dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
     _ai_edge_converter_flags: Optional[dict[str, Any]] = None,
@@ -217,6 +225,11 @@ def convert(
       prior to conversion.
     sample_kwargs: Dict of str to tensor by which the torch module will be
       traced with prior to conversion.
+    strict_export: Experimental `strict` arg for torch.export.export. When
+      enabled, the export function will trace the program through TorchDynamo
+      and ensure the soundness of the exported graph. When strict_export="auto",
+      the function will try to export module in both modes and use the first one
+      succeeds for downstream conversion.
     quant_config: User-defined quantization method and scheme of the model.
     dynamic_shapes: Optional dict or tuple that specify dynamic shape
       specifications for each input in original order. See
@@ -242,6 +255,7 @@ def convert(
       module,
       sample_args,
       sample_kwargs,
+      strict_export=strict_export,
       quant_config=quant_config,
       dynamic_shapes=dynamic_shapes,
       _ai_edge_converter_flags=_ai_edge_converter_flags,

ai_edge_torch/fx_pass_base.py

@@ -95,6 +95,15 @@ class CanonicalizePass(ExportedProgramPassBase):
   }
 
   def call(self, exported_program: torch.export.ExportedProgram):
+    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, size)
+
     exported_program = exported_program.run_decompositions(
         self._DUMMY_DECOMP_TABLE
     )

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241010.dist-info → ai_edge_torch_nightly-0.3.0.dev20241011.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20241010
+Version: 0.3.0.dev20241011
 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.3.0.dev20241010.dist-info → ai_edge_torch_nightly-0.3.0.dev20241011.dist-info}/RECORD

@@ -1,13 +1,13 @@
 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=D86Gw3pIRcpnTebUPKlnPbPGJae1S6Fw4DZZ3ZkD0zw,3730
+ai_edge_torch/fx_pass_base.py,sha256=ToyPo7MBRUNSZivwfLgAS-CNS5SnlhmZ7i_Pdzv9IJw,4238
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=bIO_IFfZmFf0EzBSomi1SB_NI_Q9y1vmOEGOd23r1GE,706
+ai_edge_torch/version.py,sha256=EdlZ-aVAW07q00kkGVtIhUqLLXc1HnI2lw1V-ymvO_0,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/conversion.py,sha256=5uPwHhmc6kwiIz-CqaiHDejf2SOWMHrb-rYEHm69wKc,3801
+ai_edge_torch/_convert/conversion.py,sha256=p1u1HJYzCxd7UyZ9_mjGOqJOZO4XeLQzmqdgO_6vL0Y,4755
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
-ai_edge_torch/_convert/converter.py,sha256=ezmaATnQi7NWDo37LUb-hEXtZSmT7_AT6vqXC6Fcq1o,8615
+ai_edge_torch/_convert/converter.py,sha256=DYbTZMZos8bvm9mLyDv3W1P8ER_iGKVohbFAmLZD4r8,9534
 ai_edge_torch/_convert/signature.py,sha256=rGpBNss3Y9FCCCcdBwDo16KqavJi8N5P0M_6Wnmoz6Y,2320
 ai_edge_torch/_convert/to_channel_last_io.py,sha256=_31phf7TYgZY2ftpNbrdlB1RhDium1lz_BXEQ6IsMFc,2893
 ai_edge_torch/_convert/fx_passes/__init__.py,sha256=xuRI0WehbUlxLHvuYjj8MeyIKBtcCp10D3E1uD1MRdw,1168
@@ -180,8 +180,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.dev20241010.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241010.dist-info/METADATA,sha256=GHeN1JauoPCMl-nkwNdqrSCv_GXFlErnAjrSVhymN9E,1897
-ai_edge_torch_nightly-0.3.0.dev20241010.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20241010.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241010.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241011.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241011.dist-info/METADATA,sha256=JlypGVS0IGIdrAWambT-SDYN2jKv63-gR0B9xL9VZLw,1897
+ai_edge_torch_nightly-0.3.0.dev20241011.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+ai_edge_torch_nightly-0.3.0.dev20241011.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241011.dist-info/RECORD,,