ai-edge-torch-nightly 0.3.0.dev20241014__py3-none-any.whl → 0.3.0.dev20241018__py3-none-any.whl

ai_edge_torch/_convert/conversion.py

@@ -35,7 +35,7 @@ def _run_convert_passes(
   exported_program = generative_fx_passes.run_generative_passes(
       exported_program
   )
-  return fx_pass_base.run_passes(
+  exported_program = fx_pass_base.run_passes(
       exported_program,
       [
           fx_passes.BuildInterpolateCompositePass(),
@@ -44,10 +44,13 @@ def _run_convert_passes(
           fx_passes.CanonicalizePass(),
           fx_passes.BuildAtenCompositePass(),
           fx_passes.CanonicalizePass(),
+          fx_passes.RemoveNonUserOutputsPass(),
+          fx_passes.CanonicalizePass(),
           fx_passes.InjectMlirDebuginfoPass(),
           fx_passes.CanonicalizePass(),
       ],
   )
+  return exported_program
 
 
 def _warn_training_modules(signatures: list[signature.Signature]):
@@ -103,17 +106,27 @@ def convert_signatures(
     nonlocal strict_export
     if strict_export == "auto":
       try:
-        return torch.export.export(*args, **kwargs, strict=True)
+        exported_program = 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)
+        exported_program = torch.export.export(*args, **kwargs, strict=False)
     elif not strict_export:
-      return torch.export.export(*args, **kwargs, strict=False)
+      exported_program = torch.export.export(*args, **kwargs, strict=False)
     else:
-      return torch.export.export(*args, **kwargs, strict=True)
+      exported_program = torch.export.export(*args, **kwargs, strict=True)
+
+    if hasattr(torch._decomp, "_decomp_table_to_post_autograd_aten"):
+      # Available after torch 2.5.0: `_decomp_table_to_post_autograd_aten` is a
+      # stop-gap table which replicates the old behaviour of post-dispatch IR.
+      # This could help ensure the collection of aten ops remaining still as the
+      # implementation of torch.export changes.
+      exported_program = exported_program.run_decompositions(
+          torch._decomp._decomp_table_to_post_autograd_aten()
+      )
+    return exported_program
 
   exported_programs: torch.export.ExportedProgram = [
       export(sig.module, sig.flat_args, dynamic_shapes=sig.dynamic_shapes)

ai_edge_torch/_convert/fx_passes/__init__.py

@@ -19,4 +19,5 @@ from ai_edge_torch._convert.fx_passes.build_aten_composite_pass import BuildAten
 from ai_edge_torch._convert.fx_passes.build_interpolate_composite_pass import BuildInterpolateCompositePass
 from ai_edge_torch._convert.fx_passes.inject_mlir_debuginfo_pass import InjectMlirDebuginfoPass
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import OptimizeLayoutTransposesPass
+from ai_edge_torch._convert.fx_passes.remove_non_user_outputs_pass import RemoveNonUserOutputsPass
 from ai_edge_torch.fx_pass_base import CanonicalizePass

ai_edge_torch/_convert/fx_passes/build_interpolate_composite_pass.py

@@ -113,6 +113,9 @@ class BuildInterpolateCompositePass(fx_pass_base.ExportedProgramPassBase):
     ]
 
   def call(self, exported_program: torch.export.ExportedProgram):
+    exported_program = fx_pass_base.run_passes(
+        exported_program, [fx_pass_base.CanonicalizePass()]
+    )
     exported_program = exported_program.run_decompositions(
         _INTERPOLATE_DECOMPOSITIONS
     )

ai_edge_torch/_convert/fx_passes/remove_non_user_outputs_pass.py

@@ -0,0 +1,52 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Pass to remove all non user outputs from exported program."""
+
+
+from ai_edge_torch import fx_pass_base
+import torch
+
+
+class RemoveNonUserOutputsPass(fx_pass_base.ExportedProgramPassBase):
+  """This pass removes all non user outputs from the exported program's output.
+
+  The FX graph may output more tensors/data than what user's original model
+  returns. Those additional outputs include user input mutations, gradient to
+  parameter, etc. Those outputs are not supported by our inference only
+  conversion or runtime. This pass remove all those outputs to ensure the
+  converted models' outputs match what returned from user's model in eval mode.
+  """
+
+  def call(self, exported_program: torch.export.ExportedProgram):
+    for node in exported_program.graph.nodes:
+      if node.op != "output":
+        continue
+
+      outputs = node.args[0]
+      output_specs = exported_program.graph_signature.output_specs
+
+      new_outputs = []
+      new_output_specs = []
+      for output, spec in zip(outputs, output_specs):
+        if spec.kind == torch.export.graph_signature.OutputKind.USER_OUTPUT:
+          new_outputs.append(output)
+          new_output_specs.append(spec)
+
+      node.args = (tuple(new_outputs),)
+      exported_program.graph_signature.output_specs = new_output_specs
+
+    exported_program.graph_module.graph.lint()
+    exported_program.graph_module.recompile()
+    return fx_pass_base.ExportedProgramPassResult(exported_program, True)

ai_edge_torch/_convert/test/test_convert.py

@@ -490,6 +490,22 @@ class TestConvert(googletest.TestCase):
     tflite_output = edge_model(**flat_inputs)
     np.testing.assert_almost_equal(reference_output, tflite_output)
 
+  def test_convert_model_with_input_mutation(self):
+    class SampleModel(nn.Module):
+
+      def forward(self, x):
+        x /= 1
+        x = x + 10
+        return x
+
+    args = (torch.randn(10, 10),)
+    torch_module = SampleModel().eval()
+    edge_model = ai_edge_torch.convert(torch_module, args)
+
+    self.assertTrue(
+        model_coverage.compare_tflite_torch(edge_model, torch_module, args)
+    )
+
 
 if __name__ == "__main__":
   googletest.main()

ai_edge_torch/fx_pass_base.py

@@ -102,7 +102,7 @@ class CanonicalizePass(ExportedProgramPassBase):
         # 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)
+        node.target = lambda self, size: torch.reshape(self.contiguous(), size)
 
     exported_program = exported_program.run_decompositions(
         self._DUMMY_DECOMP_TABLE

ai_edge_torch/generative/layers/attention_utils.py

@@ -107,7 +107,7 @@ def build_sliding_window_mask_cache(
   sliding_mask = torch.triu(all_ones, -1 * window_size + 1) * torch.tril(
       all_ones, window_size - 1
   )
-  return torch.where(sliding_mask == 1, mask, -2.3819763e38)
+  return torch.where(sliding_mask == 1, mask, float('-inf'))
 
 
 def relative_position_bucket(

ai_edge_torch/hlfb/mark_pattern/pattern.py

@@ -17,6 +17,7 @@
 import dataclasses
 from typing import Any, Callable, Optional, Union
 
+from ai_edge_torch import fx_pass_base
 from ai_edge_torch.hlfb.mark_pattern import passes
 import torch
 from torch.export.graph_signature import TensorArgument
@@ -116,6 +117,7 @@ def _find_scalar_attr(
     track_args[tracker.pattern_arg_pos] = source
     ep = torch.export.export(pattern_module, tuple(track_args))
     if decomp_table is not None:
+      ep = fx_pass_base.run_passes(ep, [fx_pass_base.CanonicalizePass()])
       ep = ep.run_decompositions(decomp_table)
 
     scalar_locs = set()
@@ -198,6 +200,9 @@ class Pattern:
 
     exported_program = torch.export.export(module, export_args)
     if decomp_table is not None:
+      exported_program = fx_pass_base.run_passes(
+          exported_program, [fx_pass_base.CanonicalizePass()]
+      )
       exported_program = exported_program.run_decompositions(decomp_table)
 
     self.exported_program = exported_program

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241014.dist-info → ai_edge_torch_nightly-0.3.0.dev20241018.dist-info}/METADATA

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

@@ -1,19 +1,20 @@
 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=ToyPo7MBRUNSZivwfLgAS-CNS5SnlhmZ7i_Pdzv9IJw,4238
+ai_edge_torch/fx_pass_base.py,sha256=SrYveglaiA_DXPoRBqSXClWM1q7853I5ujRorq_MV0M,4251
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=lCXRgmm48WXivUhIJHdj4a9DrbLbyywsh6G4yWE1Bro,706
+ai_edge_torch/version.py,sha256=K5OPFteQXEGRnfR-yFeyGvU5qSbFmO67wiOGgKTvmi0,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/conversion.py,sha256=p1u1HJYzCxd7UyZ9_mjGOqJOZO4XeLQzmqdgO_6vL0Y,4755
+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
 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
+ai_edge_torch/_convert/fx_passes/__init__.py,sha256=NVe-eGcm7j8jZpP2pcMhC8j5dVjgR1pPzyXhHdvKH4E,1267
 ai_edge_torch/_convert/fx_passes/build_aten_composite_pass.py,sha256=doaww8KqrgRTD5LotBVAIRFsEqzPn9R5lcGehBJOczA,9098
-ai_edge_torch/_convert/fx_passes/build_interpolate_composite_pass.py,sha256=vd2phg5j3Exn6BuGpASe5cU_wY4JV_YcNTssM6Q9k2c,4169
+ai_edge_torch/_convert/fx_passes/build_interpolate_composite_pass.py,sha256=XCVqWg_ask0Kb64PED0ZGAODsUuIgfyO2ZJM6aK-TXI,4283
 ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=IlZuK42kfVcRqAWZp4j2k_81T2uWo9T2558U_GPJAlU,2327
+ai_edge_torch/_convert/fx_passes/remove_non_user_outputs_pass.py,sha256=f1IUVWyhioOClsMiZzLyynoW2R17U83vA-7Q-3pGPM4,2126
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/__init__.py,sha256=lxnoH-WGLeiQIF8XjMGodjiZEFTxucl7g05N7MR9OPk,796
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_check.py,sha256=hDsl9AHzmyuSWsdHOSO114l4nBUgUdAOUWafMTipMgA,7629
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_mark.py,sha256=4RyGUwR22bZqkn_TnptenFJodc_Q43f4_SBG7gmTbos,1621
@@ -25,7 +26,7 @@ ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitio
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py,sha256=L_x8BrF7UDah-SYl-pG11I6CIckdU9kBTUHcmwW4cts,2420
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py,sha256=mzfL9cf0qBnpmxM_OlMQFvQsEZV2B_Mia9yEJV4J7rI,7135
 ai_edge_torch/_convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/test/test_convert.py,sha256=40QRxQFNeSRr4dLXJkzG-wKUlvJtsfv62cdvRrmBv5w,15097
+ai_edge_torch/_convert/test/test_convert.py,sha256=yXfeWDw9u_rTS3B6kvvFPo5E4XNT3zKTSLFSBSAI9Fc,15502
 ai_edge_torch/_convert/test/test_convert_composites.py,sha256=BCIODgxMI_3MxMLfNWYMGjcz-al-J3z5eDHCiZJXNwY,7992
 ai_edge_torch/_convert/test/test_convert_multisig.py,sha256=6_C2R9--KyNR7_oezZIAfyTSR97tOeEWy4XGcbSxBDE,5778
 ai_edge_torch/_convert/test/test_to_channel_last_io.py,sha256=1o-gUiwzIuO67FNAJ8DeyKv8fVUeZVNNNwofNVDjYeU,3024
@@ -100,7 +101,7 @@ ai_edge_torch/generative/fx_passes/__init__.py,sha256=jrzCB3ZyY_t5jJM1e2Czdt3DjA
 ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=hhxSQvkDMv0isZJhmuLiod66ZODaJ8uSPSVTJVHBabQ,1931
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/attention.py,sha256=Z0Y_G8IG0LmvLX2u9D8__Fkr22szB-az6wMNnZpzhkA,13233
-ai_edge_torch/generative/layers/attention_utils.py,sha256=68GXGR2HSWBFViTxX7cHifzVG-kcLS2IL2tQJPIpupg,7344
+ai_edge_torch/generative/layers/attention_utils.py,sha256=zBVwlBUTs-nStIKCZG0ks5ra7tsqc9ShfakFJKH5rds,7344
 ai_edge_torch/generative/layers/builder.py,sha256=XyZS1RrnMbvypeLMfwU7h1Y4x5r4WGgOx2YGJF0OUNQ,5064
 ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-mynllPiqLdXQMF6JMnc,4189
 ai_edge_torch/generative/layers/kv_cache.py,sha256=2El7kZYnQRCRcVc63xgiAdBh9oVOksDu35p9XggvaGE,6148
@@ -137,7 +138,7 @@ ai_edge_torch/generative/utilities/verifier.py,sha256=wQ4EtIED_a6FRsaOXeoQVZiHNx
 ai_edge_torch/hlfb/__init__.py,sha256=sH4um75na-O8tzxN6chFyp6Y4xnexsE7kUQpZySv6dE,735
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=cjTprggj_cuktSCm7-A25e7Shop3k63ylp7sdZmtZ8o,4790
 ai_edge_torch/hlfb/mark_pattern/passes.py,sha256=pjkKcI1nHECPluAt87cFBrt1DP0f3ge7rHq1NhCkBIE,1936
-ai_edge_torch/hlfb/mark_pattern/pattern.py,sha256=uiYRfzD1T8deCEAGfdAFusRbI41m14zeTt0Lz5lNT3M,9808
+ai_edge_torch/hlfb/mark_pattern/pattern.py,sha256=NP2mYhe5D2GjtqQfqqldp-ko3xtNghuFKKJOQskUJFI,10041
 ai_edge_torch/hlfb/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/hlfb/test/test_mark_pattern.py,sha256=ivq0eVjuf31idfNY0E12F4FxdkSI9hwYXapLJBkIf8Q,4831
 ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py,sha256=j8WpeS-mz3Zr4I7p7NwanQzkQNeH0asZ7lz5y7twgQ4,8447
@@ -180,8 +181,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.dev20241014.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241014.dist-info/METADATA,sha256=LQb4o0OMij14VvLSvUCYScCnWxxlTAp_MjNtaYI-nd4,1897
-ai_edge_torch_nightly-0.3.0.dev20241014.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20241014.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241014.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241018.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241018.dist-info/METADATA,sha256=9jDElmb5TsG683OMDoH2ywytzqOAGNM5ImTPKEKyeFc,1897
+ai_edge_torch_nightly-0.3.0.dev20241018.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+ai_edge_torch_nightly-0.3.0.dev20241018.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241018.dist-info/RECORD,,