ai-edge-torch-nightly 0.2.0.dev20240626__py3-none-any.whl → 0.2.0.dev20240628__py3-none-any.whl

ai_edge_torch/__init__.py

@@ -15,6 +15,7 @@
 
 from .convert.converter import convert
 from .convert.converter import signature
+from .convert.to_channel_last_io import to_channel_last_io
 from .model import Model
 
 

ai_edge_torch/convert/fx_passes/build_interpolate_composite_pass.py

@@ -17,10 +17,22 @@ import functools
 
 import torch
 
-from ai_edge_torch.convert.fx_passes import FxPassBase
-from ai_edge_torch.convert.fx_passes import FxPassResult
+from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassBase
+from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassResult  # NOQA
 from ai_edge_torch.hlfb import mark_pattern
 
+# For torch nightly released after mid June 2024,
+# torch.nn.functional.interpolate no longer gets exported into decomposed graph
+# but single aten op torch.ops.aten.upsample_nearest2d.vec/torch.ops.aten.upsample_bilinear2d.vec.
+# This behavior would our pattern matching based composite builder.
+# It requires the pattern and model graph to get decomposed first for backward compatibility.
+_INTERPOLATE_DECOMPOSITIONS = torch._decomp.get_decompositions(
+    [
+        torch.ops.aten.upsample_bilinear2d.vec,
+        torch.ops.aten.upsample_nearest2d.vec,
+    ]
+)
+
 
 @functools.cache
 def _get_upsample_bilinear2d_pattern():
@@ -30,6 +42,7 @@ def _get_upsample_bilinear2d_pattern():
           x, scale_factor=2, mode="bilinear", align_corners=False
       ),
       export_args=(torch.rand(1, 3, 100, 100),),
+      decomp_table=_INTERPOLATE_DECOMPOSITIONS,
   )
 
   @pattern.register_attr_builder
@@ -52,6 +65,7 @@ def _get_upsample_bilinear2d_align_corners_pattern():
           x, scale_factor=2, mode="bilinear", align_corners=True
       ),
       export_args=(torch.rand(1, 3, 100, 100),),
+      decomp_table=_INTERPOLATE_DECOMPOSITIONS,
   )
 
   @pattern.register_attr_builder
@@ -72,6 +86,7 @@ def _get_interpolate_nearest2d_pattern():
       "tfl.resize_nearest_neighbor",
       lambda x: torch.nn.functional.interpolate(x, scale_factor=2, mode="nearest"),
       export_args=(torch.rand(1, 3, 100, 100),),
+      decomp_table=_INTERPOLATE_DECOMPOSITIONS,
   )
 
   @pattern.register_attr_builder
@@ -86,7 +101,7 @@ def _get_interpolate_nearest2d_pattern():
   return pattern
 
 
-class BuildInterpolateCompositePass(FxPassBase):
+class BuildInterpolateCompositePass(ExportedProgramPassBase):
 
   def __init__(self):
     super().__init__()
@@ -96,10 +111,13 @@ class BuildInterpolateCompositePass(FxPassBase):
         _get_interpolate_nearest2d_pattern(),
     ]
 
-  def call(self, graph_module: torch.fx.GraphModule):
+  def call(self, exported_program: torch.export.ExportedProgram):
+    exported_program = exported_program.run_decompositions(_INTERPOLATE_DECOMPOSITIONS)
+
+    graph_module = exported_program.graph_module
     for pattern in self._patterns:
       graph_module = mark_pattern.mark_pattern(graph_module, pattern)
 
     graph_module.graph.lint()
     graph_module.recompile()
-    return FxPassResult(graph_module, True)
+    return ExportedProgramPassResult(exported_program, True)

ai_edge_torch/convert/test/test_to_channel_last_io.py

@@ -0,0 +1,96 @@
+# 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.
+# ==============================================================================
+
+import unittest
+
+import torch
+
+import ai_edge_torch
+
+
+class Identity(torch.nn.Module):
+
+  def forward(self, x):
+    return x
+
+
+class TestToChannelLastIO(unittest.TestCase):
+  """Tests to_channel_last_io API and module wrapper."""
+
+  def test_no_transformations(self):
+    x = torch.rand(1, 3, 10, 10)
+    y = ai_edge_torch.to_channel_last_io(Identity())(x)
+    self.assertEqual(y.shape, (1, 3, 10, 10))
+
+  def test_args(self):
+    x = torch.rand(1, 10, 10, 3)
+    y = ai_edge_torch.to_channel_last_io(Identity(), args=[0])(x)
+    self.assertEqual(y.shape, (1, 3, 10, 10))
+
+  def test_outputs(self):
+    x = torch.rand(1, 3, 10, 10)
+    y = ai_edge_torch.to_channel_last_io(Identity(), outputs=[0])(x)
+    self.assertEqual(y.shape, (1, 10, 10, 3))
+
+  def test_args_outputs(self):
+    x = torch.rand(1, 10, 10, 3)
+    y = ai_edge_torch.to_channel_last_io(Identity(), args=[0], outputs=[0])(x)
+    self.assertEqual(y.shape, (1, 10, 10, 3))
+
+  def test_args_5d(self):
+    x = torch.rand(1, 10, 10, 10, 3)
+    y = ai_edge_torch.to_channel_last_io(Identity(), args=[0])(x)
+    self.assertEqual(y.shape, (1, 3, 10, 10, 10))
+
+  def test_outputs_5d(self):
+    x = torch.rand(1, 3, 10, 10, 10)
+    y = ai_edge_torch.to_channel_last_io(Identity(), outputs=[0])(x)
+    self.assertEqual(y.shape, (1, 10, 10, 10, 3))
+
+  def test_chained_wrappers(self):
+    x = torch.rand(1, 10, 10, 3)
+
+    m = Identity()
+    m = ai_edge_torch.to_channel_last_io(m, args=[0])
+    m = ai_edge_torch.to_channel_last_io(m, outputs=[0])
+
+    y = m(x)
+    self.assertEqual(y.shape, (1, 10, 10, 3))
+
+  def test_list_args(self):
+    class Add(torch.nn.Module):
+
+      def forward(self, x, y):
+        return x + y
+
+    x = (torch.rand(1, 10, 10, 3), torch.rand(1, 10, 10, 3))
+    y = ai_edge_torch.to_channel_last_io(Add(), args=[0, 1])(*x)
+    self.assertEqual(y.shape, (1, 3, 10, 10))
+
+  def test_list_outputs(self):
+    class TwoIdentity(torch.nn.Module):
+
+      def forward(self, x):
+        return x, x
+
+    x = torch.rand(1, 3, 10, 10)
+    y = ai_edge_torch.to_channel_last_io(TwoIdentity(), outputs=[0])(x)
+    self.assertIsInstance(y, tuple)
+    self.assertEqual(y[0].shape, (1, 10, 10, 3))
+    self.assertEqual(y[1].shape, (1, 3, 10, 10))
+
+
+if __name__ == "__main__":
+  unittest.main()

ai_edge_torch/convert/to_channel_last_io.py

@@ -0,0 +1,85 @@
+# 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.
+# ==============================================================================
+
+from typing import Optional
+
+import torch
+from torch import nn
+
+
+class ChannelLastIOWrapper(nn.Module):
+
+  def __init__(self, wrapped, *, args=None, outputs=None):
+    super().__init__()
+    self.wrapped = wrapped
+    self._args = args or []
+    self._outputs = outputs or []
+
+  def _to_channel_last(self, x):
+    if not torch.is_tensor(x):
+      raise ValueError("Input must be a torch tensor")
+    if x.ndim < 3:
+      raise ValueError("Input must be a tensor with rank >= 3 in layout (N, C, ...)")
+    dims = [0, *range(2, x.ndim), 1]
+    return torch.permute(x, dims)
+
+  def _to_channel_first(self, x):
+    if not torch.is_tensor(x):
+      raise ValueError("Input must be a torch tensor.")
+    if x.ndim < 3:
+      raise ValueError("Input must be a tensor with rank >= 3 in layout (N, ..., C)")
+    dims = [0, x.ndim - 1, *range(1, x.ndim - 1)]
+    return torch.permute(x, dims)
+
+  def forward(self, *args, **kwargs):
+    args = list(args)
+    for i in self._args:
+      args[i] = self._to_channel_first(args[i])
+
+    outputs = self.wrapped(*args, **kwargs)
+
+    if not isinstance(outputs, (list, tuple)):
+      outputs_is_list = False
+      output_list = [outputs]
+    else:
+      outputs_is_list = True
+      output_list = list(outputs)
+
+    for i in self._outputs:
+      output_list[i] = self._to_channel_last(output_list[i])
+
+    if not outputs_is_list:
+      return output_list[0]
+    else:
+      return type(outputs)(output_list)
+
+
+def to_channel_last_io(
+    module: nn.Module,
+    args: Optional[list[int]] = None,
+    outputs: Optional[list[int]] = None,
+):
+  """Wraps the module with channel first to channel last layout transformations.
+
+  Args:
+    args (list[int]): Transform args with indices in the list from channel first
+      (N, C, ...) to channel last (N, ..., C).
+    outputs (list[int]): Transform outputs with indices in the list from channel
+      first (N, C, ...) to channel last (N, ..., C).
+  Returns:
+    The wrapped nn.Module with additional layout transposes after inputs and/or before
+    outputs.
+  """
+  return ChannelLastIOWrapper(module, args=args, outputs=outputs)

ai_edge_torch/hlfb/mark_pattern/pattern.py

@@ -100,13 +100,18 @@ class ScalarAttrLocation:
 
 
 def _find_scalar_attr(
-    pattern_module: torch.nn.Module, export_args: tuple[Any], tracker: ScalarAttrTracker
+    pattern_module: torch.nn.Module,
+    export_args: tuple[Any],
+    tracker: ScalarAttrTracker,
+    decomp_table=None,
 ) -> ScalarAttrLocation:
   scalar_loc_intersections = None
   for source, target in tracker._source_targets:
     track_args = list(export_args)
     track_args[tracker.pattern_arg_pos] = source
     ep = torch.export.export(pattern_module, tuple(track_args))
+    if decomp_table is not None:
+      ep = ep.run_decompositions(decomp_table)
 
     scalar_locs = set()
     nodes = ep.graph_module.graph.nodes
@@ -145,6 +150,7 @@ class Pattern:
           ["Pattern", GraphModule, InternalMatch], Optional[dict[str, Any]]
       ] = None,
       scalar_attr_trackers: list[ScalarAttrTracker] = None,
+      decomp_table: Optional[dict[torch._ops.OperatorBase, Callable]] = None,
   ):
     """The PyTorch computation pattern to match against a model.
 
@@ -165,6 +171,8 @@ class Pattern:
         for scalar args in `export_args`, which are used to track
         the attr occurrence(s) and retrieve their values from the
         matched subgraph.
+      decomp_table (Optional[dict[torch._ops.OperatorBase, Callable]]):
+        The decomposition table to be run on the pattern's exported program.
     """
     if not isinstance(module, torch.nn.Module):
 
@@ -180,23 +188,28 @@ class Pattern:
       module = PatternModule(module).eval()
 
     self.name = name
-    self.exported_program = torch.export.export(module, export_args)
-    self.graph_module = self.exported_program.graph_module
     self.attr_builder = attr_builder
     self._scalar_attr_trackers = scalar_attr_trackers if scalar_attr_trackers else []
 
-    # Sanitize graph_module for more precise pattern matching.
-    # The graph_module to match against this pattern should apply equivalent
-    # sanitization.
-    self.graph_module = passes.remove_clone_ops(self.graph_module)
-    self.graph_module = passes.remove_dangling_args(self.graph_module)
+    exported_program = torch.export.export(module, export_args)
+    if decomp_table is not None:
+      exported_program = exported_program.run_decompositions(decomp_table)
+
+    self.exported_program = exported_program
+    self.graph_module = self.exported_program.graph_module
 
     self._scalar_attr_locations = []
     for tracker in self._scalar_attr_trackers:
       self._scalar_attr_locations.append(
-          _find_scalar_attr(module, export_args, tracker)
+          _find_scalar_attr(module, export_args, tracker, decomp_table=decomp_table)
       )
 
+    # Sanitize graph_module for more precise pattern matching.
+    # The graph_module to match against this pattern should apply equivalent
+    # sanitization.
+    self.graph_module = passes.remove_clone_ops(self.graph_module)
+    self.graph_module = passes.remove_dangling_args(self.graph_module)
+
     # Builds list of ordered input and output nodes.
     self.graph_nodes_map = {}
     for node in self.graph_module.graph.nodes:

{ai_edge_torch_nightly-0.2.0.dev20240626.dist-info → ai_edge_torch_nightly-0.2.0.dev20240628.dist-info}/METADATA

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

@@ -1,13 +1,14 @@
-ai_edge_torch/__init__.py,sha256=FPMmuFU3pyMREtjB_san1fy_0PFtAsgA0VZfOYvDrb4,1008
+ai_edge_torch/__init__.py,sha256=CNDboRP4zQBpz2hznNCQWcQCARvNXUm3DMa1Dw_XXFg,1067
 ai_edge_torch/model.py,sha256=kmcgELjsYl8YzF8nUF6P7q4i8MWS-pLGpfsy-yTUXmE,4243
 ai_edge_torch/convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/convert/conversion.py,sha256=8K8jQuaCjlUWoj7jiimxp_zpN6mYThLOcQ858UDcYnE,4159
 ai_edge_torch/convert/conversion_utils.py,sha256=9BqCL38DErv1vEVGtT3BIJVhdwZjw2EQ-_m5UpvVVYE,11294
 ai_edge_torch/convert/converter.py,sha256=bjj5TV5_g4sGyuSh8ThEDydlNMqhkGSY4SzXK6vwhqI,6927
+ai_edge_torch/convert/to_channel_last_io.py,sha256=zo5tY3yDhY_EPCkrL1XSXs2uRFS8B4_qu08dSjNsUGk,2778
 ai_edge_torch/convert/fx_passes/__init__.py,sha256=EPs4PSIDLuRH5EBETi6deaOvaaf_Q4xD3_9NVcR7x8o,2810
 ai_edge_torch/convert/fx_passes/_pass_base.py,sha256=ijVyDclPnd6a0DWWUJkwR4igj6f82S-cE1-83QGPvgw,1652
 ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py,sha256=2yqUwJJ2R233_X9FNMOP9oYRTTzH34TR_BIUj-wfnKw,7080
-ai_edge_torch/convert/fx_passes/build_interpolate_composite_pass.py,sha256=jB27GlDC8x36nn35aiq97uKERiq4KXSUZ7tv7yc0Gl4,3223
+ai_edge_torch/convert/fx_passes/build_interpolate_composite_pass.py,sha256=6m_vcycd9f3OQgQLx2hhQjsKfOqdxE5EkjzqrxqyAQM,4168
 ai_edge_torch/convert/fx_passes/canonicalize_pass.py,sha256=UX6dJsxCqSkftXXvNBV-i7Bjk6H7qTyqzUnE640Itfg,1673
 ai_edge_torch/convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=aRT8hTS3n9ie28lgu6mygtFO6Ypwu0qjNb0c81v9HLs,2448
 ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/__init__.py,sha256=VA9bekxPVhLk4MYlIRXnOzrSnbCtUmGj7OQ_fJcKQtc,795
@@ -24,6 +25,7 @@ ai_edge_torch/convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrK
 ai_edge_torch/convert/test/test_convert.py,sha256=2qPmmGqnfV_o1gfsSdjGq3-JR1b323ligiy5MdAv9NA,8021
 ai_edge_torch/convert/test/test_convert_composites.py,sha256=_Ojc-H6GOS5s8ek3_8eRBL_AiCs-k3srziPJ2R4Ulrg,7255
 ai_edge_torch/convert/test/test_convert_multisig.py,sha256=kMaGnHe9ylfyU68qCifYcaGwJqyejKz--QQt9jS2oUA,4537
+ai_edge_torch/convert/test/test_to_channel_last_io.py,sha256=I8c4ZG3v1vo0yxQYzLK_BTId4AOL9vadHGDtfCUZ4UI,2930
 ai_edge_torch/debug/__init__.py,sha256=N05Mmvi41KgSuK0JhuMejERESgP8QekiGdp9_PEyuKU,742
 ai_edge_torch/debug/culprit.py,sha256=urtCKPXORPvn6oyDxDSCSjgvngUnjjcsUMwAOeIl15E,14236
 ai_edge_torch/debug/utils.py,sha256=hjVmQVVl1dKxEF0D6KB4a3ouQ3wBkTsebOX2YsUObZM,1430
@@ -101,7 +103,7 @@ ai_edge_torch/generative/utilities/t5_loader.py,sha256=h1FQzt4x8wiQMX4NzYNVIaJGL
 ai_edge_torch/hlfb/__init__.py,sha256=rrje8a2iuKboBoV96bVq7nlS9HsnuEMbHE5JiWmCxFA,752
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=2VXnHcGf23VOuP-1GriGIpuL98leBB8twp_qaScMnmc,4799
 ai_edge_torch/hlfb/mark_pattern/passes.py,sha256=YV2YKBkh7y7j7sd7EA81vf_1hUKUvTRiy1pfqZustXc,1539
-ai_edge_torch/hlfb/mark_pattern/pattern.py,sha256=H4047w-xwx27rYPKNqmeOSQ9M1Adkpd7drp81YdV7Hw,9206
+ai_edge_torch/hlfb/mark_pattern/pattern.py,sha256=BEMKgkX8IrsX70h2CxwA_tpsBm6qWWe2bOeOufMYNkA,9722
 ai_edge_torch/hlfb/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/hlfb/test/test_mark_pattern.py,sha256=qYR3PRGS9W3OG-qX7UFqL69VxXuUSfyDBUJtCXtXcOE,4262
 ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py,sha256=aUAPKnH4_Jxpp3mLlD5rzdT1g_VBm7OrwwLJ9DeJlzQ,8190
@@ -112,8 +114,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=eO9Ra160ITjQSyRBEGy6nNIVH3gYacSWDd
 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=EIyKz-HY70DguWuSrJal8LpYXQ5ZSEUf3ZrVl7jikFM,4286
-ai_edge_torch_nightly-0.2.0.dev20240626.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240626.dist-info/METADATA,sha256=taaGTe-WdFG7HWUB43xnjYmG_iFcRwaIaZCK91-QP7M,1748
-ai_edge_torch_nightly-0.2.0.dev20240626.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240626.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240626.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240628.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240628.dist-info/METADATA,sha256=Jt1FVBxKIu2RO5oK2Y3bdu-IuknkTyxxRxa9VXMMUjA,1748
+ai_edge_torch_nightly-0.2.0.dev20240628.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240628.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240628.dist-info/RECORD,,