ai-edge-torch-nightly 0.3.0.dev20250121__py3-none-any.whl → 0.3.0.dev20250123__py3-none-any.whl

ai_edge_torch/__init__.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
 from ai_edge_torch._config import config
 from ai_edge_torch._convert.converter import convert
 from ai_edge_torch._convert.converter import signature
@@ -20,6 +19,7 @@ from ai_edge_torch._convert.to_channel_last_io import to_channel_last_io
 from ai_edge_torch.model import Model
 from ai_edge_torch.version import __version__
 
+
 def load(path: str) -> Model:
   """Imports an ai_edge_torch model from disk.
 

ai_edge_torch/_convert/conversion.py

@@ -17,7 +17,7 @@ import logging
 from typing import Any, Literal, Optional, Union
 
 import ai_edge_torch
-from ai_edge_torch import fx_pass_base
+from ai_edge_torch import fx_infra
 from ai_edge_torch import lowertools
 from ai_edge_torch import model
 from ai_edge_torch._convert import fx_passes
@@ -53,7 +53,7 @@ def _run_convert_passes(
         fx_passes.CanonicalizePass(),
     ]
 
-  exported_program = fx_pass_base.run_passes(exported_program, passes)
+  exported_program = fx_infra.run_passes(exported_program, passes)
   return exported_program
 
 
@@ -125,14 +125,10 @@ def convert_signatures(
     else:
       exported_program = torch.export.export(**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()
-      )
+    exported_program = fx_infra.safe_run_decompositions(
+        exported_program,
+        fx_infra.decomp.pre_convert_decomp(),
+    )
     return exported_program
 
   exported_programs: torch.export.ExportedProgram = [

ai_edge_torch/_convert/fx_passes/__init__.py

@@ -20,4 +20,4 @@ from ai_edge_torch._convert.fx_passes.build_interpolate_composite_pass import Bu
 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
+from ai_edge_torch.fx_infra import CanonicalizePass

ai_edge_torch/_convert/fx_passes/build_aten_composite_pass.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 
 from typing import Any, Callable
-from ai_edge_torch import fx_pass_base
+from ai_edge_torch import fx_infra
 from ai_edge_torch import lowertools
 import torch
 import torch.utils._pytree as pytree
@@ -25,6 +25,9 @@ _composite_builders: dict[
 
 
 def _register_composite_builder(op):
+  # Remove op from pre_convert_decomp to keep this in the decomposed graph.
+  fx_infra.decomp.remove_pre_convert_decomp(op)
+
   def inner(func):
     if isinstance(op, torch._ops.OpOverloadPacket):
       for overload in op.overloads():
@@ -276,7 +279,7 @@ def _aten_embedding(gm: torch.fx.GraphModule, node: torch.fx.Node):
   node.target = embedding
 
 
-class BuildAtenCompositePass(fx_pass_base.PassBase):
+class BuildAtenCompositePass(fx_infra.PassBase):
 
   def call(self, graph_module: torch.fx.GraphModule):
     for node in graph_module.graph.nodes:
@@ -285,4 +288,4 @@ class BuildAtenCompositePass(fx_pass_base.PassBase):
 
     graph_module.graph.lint()
     graph_module.recompile()
-    return fx_pass_base.PassResult(graph_module, True)
+    return fx_infra.PassResult(graph_module, True)

ai_edge_torch/_convert/fx_passes/build_interpolate_composite_pass.py

@@ -16,7 +16,7 @@
 
 import functools
 
-from ai_edge_torch import fx_pass_base
+from ai_edge_torch import fx_infra
 from ai_edge_torch.hlfb import mark_pattern
 from ai_edge_torch.hlfb.mark_pattern import pattern as pattern_module
 import torch
@@ -41,7 +41,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,
+      extra_decomp_table=_INTERPOLATE_DECOMPOSITIONS,
   )
 
   @pattern.register_attr_builder
@@ -65,7 +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,
+      extra_decomp_table=_INTERPOLATE_DECOMPOSITIONS,
   )
 
   @pattern.register_attr_builder
@@ -89,7 +89,7 @@ def _get_interpolate_nearest2d_pattern():
           x, scale_factor=2, mode="nearest"
       ),
       export_args=(torch.rand(1, 3, 100, 100),),
-      decomp_table=_INTERPOLATE_DECOMPOSITIONS,
+      extra_decomp_table=_INTERPOLATE_DECOMPOSITIONS,
   )
 
   @pattern.register_attr_builder
@@ -104,7 +104,7 @@ def _get_interpolate_nearest2d_pattern():
   return pattern
 
 
-class BuildInterpolateCompositePass(fx_pass_base.ExportedProgramPassBase):
+class BuildInterpolateCompositePass(fx_infra.ExportedProgramPassBase):
 
   def __init__(self):
     super().__init__()
@@ -115,11 +115,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
+    exported_program = fx_infra.safe_run_decompositions(
+        exported_program,
+        _INTERPOLATE_DECOMPOSITIONS,
     )
 
     graph_module = exported_program.graph_module
@@ -128,4 +126,4 @@ class BuildInterpolateCompositePass(fx_pass_base.ExportedProgramPassBase):
 
     graph_module.graph.lint()
     graph_module.recompile()
-    return fx_pass_base.ExportedProgramPassResult(exported_program, True)
+    return fx_infra.ExportedProgramPassResult(exported_program, True)

ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from ai_edge_torch import fx_pass_base
+from ai_edge_torch import fx_infra
 from ai_edge_torch import lowertools
 import torch
 import torch.utils._pytree as pytree
@@ -61,7 +61,7 @@ def _wrap_call_function_node_with_debuginfo_writer(node: torch.fx.GraphModule):
   node.target = debuginfo_writer
 
 
-class InjectMlirDebuginfoPass(fx_pass_base.PassBase):
+class InjectMlirDebuginfoPass(fx_infra.PassBase):
   """DEPRECATED: Debuginfo is injected automatically by odml_torch."""
 
   def call(self, graph_module: torch.fx.GraphModule):
@@ -70,4 +70,4 @@ class InjectMlirDebuginfoPass(fx_pass_base.PassBase):
 
     graph_module.graph.lint()
     graph_module.recompile()
-    return fx_pass_base.PassResult(graph_module, True)
+    return fx_infra.PassResult(graph_module, True)

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/__init__.py

@@ -13,4 +13,5 @@
 # limitations under the License.
 # ==============================================================================
 
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import _decomp_registry
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass.pass_body import OptimizeLayoutTransposesPass  # NOQA

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/_decomp_registry.py

@@ -0,0 +1,33 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Remove decompositions for ops to keep in layout optimization."""
+from ai_edge_torch import fx_infra
+import torch
+
+__all__ = []
+
+aten = torch.ops.aten
+
+_OPS_TO_KEEP = [
+    aten.conv2d,
+    aten.max_pool2d,
+    aten._softmax.default,
+    aten.group_norm.default,
+    aten.native_group_norm.default,
+    aten.reflection_pad2d.default,
+]
+
+for op in _OPS_TO_KEEP:
+  fx_infra.decomp.remove_pre_convert_decomp(op)

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/op_func_registry.py

@@ -20,6 +20,7 @@ from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import uti
 class OpFuncRegistry(dict):
 
   def register(self, op):
+
     ops = utils.flatten_torch_op_overloads(op)
 
     def inner(func):

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/pass_body.py

@@ -18,7 +18,7 @@ import operator
 import os
 from typing import Union
 
-from ai_edge_torch import fx_pass_base
+from ai_edge_torch import fx_infra
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_check  # NOQA
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_mark  # NOQA
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_partitioners  # NOQA
@@ -30,7 +30,7 @@ import torch.ao.quantization.quantize_pt2e
 TransposeFunc = Union[utils.tensor_to_nchw, utils.tensor_to_nhwc]
 
 
-class OptimizeLayoutTransposesPass(fx_pass_base.ExportedProgramPassBase):
+class OptimizeLayoutTransposesPass(fx_infra.ExportedProgramPassBase):
 
   def get_source_meta(self, node: torch.fx.Node):
     keys = ["stack_trace", "nn_module_stack", "source_fn_stack", "from_node"]
@@ -300,4 +300,4 @@ class OptimizeLayoutTransposesPass(fx_pass_base.ExportedProgramPassBase):
     # Mark const node again for debugging
     self.mark_const_nodes(exported_program)
 
-    return fx_pass_base.ExportedProgramPassResult(exported_program, True)
+    return fx_infra.ExportedProgramPassResult(exported_program, True)

ai_edge_torch/_convert/fx_passes/remove_non_user_outputs_pass.py

@@ -15,11 +15,11 @@
 """Pass to remove all non user outputs from exported program."""
 
 
-from ai_edge_torch import fx_pass_base
+from ai_edge_torch import fx_infra
 import torch
 
 
-class RemoveNonUserOutputsPass(fx_pass_base.ExportedProgramPassBase):
+class RemoveNonUserOutputsPass(fx_infra.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
@@ -47,6 +47,6 @@ class RemoveNonUserOutputsPass(fx_pass_base.ExportedProgramPassBase):
       node.args = (tuple(new_outputs),)
       exported_program.graph_signature.output_specs = new_output_specs
 
-    exported_program.graph_module.graph.lint()
+    exported_program.graph.eliminate_dead_code()
     exported_program.graph_module.recompile()
-    return fx_pass_base.ExportedProgramPassResult(exported_program, True)
+    return fx_infra.ExportedProgramPassResult(exported_program, True)

ai_edge_torch/_convert/test/test_convert.py

@@ -511,7 +511,7 @@ class TestConvert(googletest.TestCase):
     # Step 1: export resnet18
     args = (torch.randn(1, 3, 224, 224),)
     m = torchvision.models.resnet18().eval()
-    m = torch._export.capture_pre_autograd_graph(m, args)
+    m = torch.export.export_for_training(m, args).module()
 
     # Step 2: Insert observers or fake quantize modules
     quantizer = pt2e_quantizer.PT2EQuantizer().set_global(
@@ -533,7 +533,7 @@ class TestConvert(googletest.TestCase):
     # Step 1: export resnet18
     args = (torch.randn(1, 3, 224, 224),)
     m = torchvision.models.resnet18().eval()
-    m = torch._export.capture_pre_autograd_graph(m, args)
+    m = torch.export.export_for_training(m, args).module()
 
     # Step 2: Insert observers or fake quantize modules
     quantizer = pt2e_quantizer.PT2EQuantizer().set_global(

ai_edge_torch/fx_infra/__init__.py

@@ -0,0 +1,32 @@
+# Copyright 2025 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 ai_edge_torch.fx_infra import _canonicalize_pass
+from ai_edge_torch.fx_infra import _safe_run_decompositions
+from ai_edge_torch.fx_infra import decomp
+from ai_edge_torch.fx_infra import graph_utils
+from ai_edge_torch.fx_infra import pass_base
+
+
+PassBase = pass_base.PassBase
+PassResult = pass_base.PassResult
+FxPassBase = pass_base.FxPassBase
+FxPassResult = pass_base.FxPassResult
+ExportedProgramPassBase = pass_base.ExportedProgramPassBase
+ExportedProgramPassResult = pass_base.ExportedProgramPassResult
+run_passes = pass_base.run_passes
+
+CanonicalizePass = _canonicalize_pass.CanonicalizePass
+safe_run_decompositions = _safe_run_decompositions.safe_run_decompositions

ai_edge_torch/fx_infra/_canonicalize_pass.py

@@ -0,0 +1,27 @@
+# Copyright 2025 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 ai_edge_torch.fx_infra import _safe_run_decompositions
+from ai_edge_torch.fx_infra import pass_base
+import torch
+
+
+class CanonicalizePass(pass_base.ExportedProgramPassBase):
+
+  def call(self, exported_program: torch.export.ExportedProgram):
+    exported_program = _safe_run_decompositions.safe_run_decompositions(
+        exported_program, {}
+    )
+
+    return pass_base.ExportedProgramPassResult(exported_program, True)

ai_edge_torch/fx_infra/_safe_run_decompositions.py

@@ -0,0 +1,57 @@
+# 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.
+# ==============================================================================
+"""ExportedProgram.run_decompositions wrapper to handle unexpected export behavior."""
+import torch
+
+
+# A dummy decomp table for running ExportedProgram.run_decompositions without
+# any op decompositions but just aot_export_module. Due to the check in
+# run_decompositions, if None or an empty dict is passed as decomp_table,
+# it will run the default aten-coreaten decompositions. Therefore a non-empty
+# dummy decomp table is needed.
+# Ref: https://github.com/pytorch/pytorch/blob/db895ace1d36726e64781774f53b3d3098206116/torch/export/exported_program.py#L543
+_DUMMY_DECOMP_TABLE = {
+    torch._ops.OperatorBase(): lambda: None,
+}
+
+
+def safe_run_decompositions(exported_program, decomp_table=None):
+  """Wrapper for ExportedProgram.run_decompositions to handle unexpected export behavior."""
+
+  if decomp_table is not None and not decomp_table:
+    # Empty decomp table means no op decompositions. Use dummy decomp table
+    # instead for backward compatibility.
+    decomp_table = _DUMMY_DECOMP_TABLE
+
+  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)
+
+  exported_program = exported_program.run_decompositions(decomp_table)
+
+  if hasattr(torch.ops.aten, "_assert_tensor_metadata"):
+    for node in exported_program.graph.nodes:
+      if node.target == torch.ops.aten._assert_tensor_metadata.default:
+        exported_program.graph.erase_node(node)
+
+  exported_program.graph.eliminate_dead_code()
+  exported_program.graph_module.recompile()
+
+  return exported_program

ai_edge_torch/fx_infra/decomp.py

@@ -0,0 +1,80 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""ExportedProgram.run_decompositions wrapper to handle unexpected export behavior."""
+import torch
+
+# Decompositions to be run after torch.export before conversion and any
+# passes. Remove ops from this decomposition table if they need to be preserved
+# in passes.
+_pre_convert_decomp = torch._decomp.core_aten_decompositions().copy()
+
+
+# Decompositions to be run after conversion before odml_torch passes and
+# lowerings.
+_pre_lower_decomp = torch._decomp.core_aten_decompositions().copy()
+
+
+def _get_ops(op):
+  if isinstance(op, torch._ops.OpOverloadPacket):
+    return [getattr(op, overload) for overload in op.overloads()]
+  else:
+    return [op]
+
+
+def pre_convert_decomp():
+  """Decompositions to be run after torch.export before conversion and any passes."""
+  return _pre_convert_decomp.copy()
+
+
+def pre_lower_decomp():
+  """Decompositions to be run after conversion before odml_torch passes and lowerings."""
+  return _pre_lower_decomp.copy()
+
+
+def remove_pre_lower_decomp(op):
+  # Also remove from pre_convert_decomp which always run before pre_lower_
+  # decomp.
+  remove_pre_convert_decomp(op)
+
+  for op_ in _get_ops(op):
+    _pre_lower_decomp.pop(op_, None)
+
+
+def remove_pre_convert_decomp(op):
+  for op_ in _get_ops(op):
+    _pre_convert_decomp.pop(op_, None)
+
+
+def add_pre_convert_decomp(op, decomp):
+  # Also add decomp to pre_lower_decomp which runs after pre_convert_decomp.
+  add_pre_lower_decomp(op, decomp)
+
+  for op_ in _get_ops(op):
+    _pre_convert_decomp[op_] = decomp
+
+
+def add_pre_lower_decomp(op, decomp):
+  for op_ in _get_ops(op):
+    _pre_lower_decomp[op_] = decomp
+
+
+def update_pre_convert_decomp(decomps):
+  for op, decomp in decomps.items():
+    add_pre_convert_decomp(op, decomp)
+
+
+def update_pre_lower_decomp(decomps):
+  for op, decomp in decomps.items():
+    add_pre_lower_decomp(op, decomp)

ai_edge_torch/fx_infra/graph_utils.py

@@ -0,0 +1,42 @@
+# 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.
+# ==============================================================================
+"""FX graph utilities."""
+import torch
+
+
+def remove_dangling_args(graph_module: torch.fx.GraphModule):
+  """Removes dangling args from the graph."""
+  for node in graph_module.graph.nodes:
+    if node.op == "placeholder" and not node.users:
+      graph_module.graph.erase_node(node)
+
+  graph_module.graph.lint()
+  graph_module.recompile()
+  return graph_module
+
+
+def remove_assert_tensor_metadata_nodes(graph_module: torch.fx.GraphModule):
+  """Removes aten._assert_tensor_metadata nodes from the graph.
+
+  This op is inserted by torch.export to check tensor metadata on custom ops. It
+  can break patten matching and lowering.
+  """
+  for node in graph_module.graph.nodes:
+    if node.target == torch.ops.aten._assert_tensor_metadata.default:
+      graph_module.graph.erase_node(node)
+
+  graph_module.graph.lint()
+  graph_module.recompile()
+  return graph_module

ai_edge_torch/{fx_pass_base.py → fx_infra/pass_base.py}

@@ -80,31 +80,3 @@ def run_passes(
             constants=exported_program.constants,
         )
   return exported_program
-
-
-class CanonicalizePass(ExportedProgramPassBase):
-
-  # A dummy decomp table for running ExportedProgram.run_decompositions without
-  # any op decompositions but just aot_export_module. Due to the check in
-  # run_decompositions, if None or an empty dict is passed as decomp_table,
-  # it will run the default aten-coreaten decompositions. Therefore a non-empty
-  # dummy decomp table is needed.
-  # Ref: https://github.com/pytorch/pytorch/blob/db895ace1d36726e64781774f53b3d3098206116/torch/export/exported_program.py#L543
-  _DUMMY_DECOMP_TABLE = {
-      torch._ops.OperatorBase(): lambda: None,
-  }
-
-  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.contiguous(), size)
-
-    exported_program = exported_program.run_decompositions(
-        self._DUMMY_DECOMP_TABLE
-    )
-    return ExportedProgramPassResult(exported_program, True)

ai_edge_torch/generative/fx_passes/__init__.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from ai_edge_torch import fx_pass_base
-from ai_edge_torch.fx_pass_base import CanonicalizePass
+from ai_edge_torch import fx_infra
+from ai_edge_torch.fx_infra import CanonicalizePass
 from ai_edge_torch.generative.fx_passes.remove_sdpa_zero_mask_pass import RemoveSDPACompositeZeroMaskPass
 import torch
 
@@ -21,7 +21,7 @@ import torch
 def run_generative_passes(
     exported_program: torch.export.ExportedProgram,
 ) -> torch.export.ExportedProgram:
-  return fx_pass_base.run_passes(
+  return fx_infra.run_passes(
       exported_program,
       [
           RemoveSDPACompositeZeroMaskPass(),

ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from ai_edge_torch import fx_pass_base
+from ai_edge_torch import fx_infra
 from ai_edge_torch import lowertools
 import torch
 
 
-class RemoveSDPACompositeZeroMaskPass(fx_pass_base.ExportedProgramPassBase):
+class RemoveSDPACompositeZeroMaskPass(fx_infra.ExportedProgramPassBase):
 
   def is_zero_tensor_node(self, node: torch.fx.Node):
     return node.target == torch.ops.aten.zeros.default
@@ -47,4 +47,4 @@ class RemoveSDPACompositeZeroMaskPass(fx_pass_base.ExportedProgramPassBase):
 
     exported_program.graph_module.graph.lint()
     exported_program.graph_module.recompile()
-    return fx_pass_base.ExportedProgramPassResult(exported_program, True)
+    return fx_infra.ExportedProgramPassResult(exported_program, True)

ai_edge_torch/hlfb/mark_pattern/__init__.py

@@ -90,7 +90,6 @@ def mark_pattern(
   graph_module_to_match = fx_utils.remove_clone_ops(graph_module_to_match)
 
   match_with_attrs = pattern.match(graph_module_to_match)
-
   for match, attr in match_with_attrs:
     match_id = _get_uuid()
 

ai_edge_torch/hlfb/mark_pattern/fx_utils.py

@@ -14,8 +14,13 @@
 # ==============================================================================
 """FX graph utilities for pattern matching clean ups."""
 
+from ai_edge_torch import fx_infra
 import torch
 
+remove_dangling_args = fx_infra.graph_utils.remove_dangling_args
+remove_assert_tensor_metadata_nodes = (
+    fx_infra.graph_utils.remove_assert_tensor_metadata_nodes
+)
 
 def is_clone_op(node: torch.fx.Node) -> bool:
   """Checks if the node is a clone op."""
@@ -46,24 +51,3 @@ def remove_clone_ops(gm: torch.fx.GraphModule):
   gm.graph.lint()
   gm.recompile()
   return gm
-
-
-def remove_dangling_args(gm: torch.fx.GraphModule):
-  """Removes dangling args from the graph.
-
-  Args:
-    gm: The graph module to remove dangling args from.
-
-  Returns:
-    The graph module with dangling args removed.
-  """
-  nodes_to_erase = []
-  for node in gm.graph.nodes:
-    if node.op == "placeholder" and len(node.users) == 0:
-      nodes_to_erase.append(node)
-  for node in nodes_to_erase:
-    gm.graph.erase_node(node)
-
-  gm.graph.lint()
-  gm.recompile()
-  return gm

ai_edge_torch/hlfb/mark_pattern/pattern.py

@@ -17,7 +17,7 @@
 import dataclasses
 from typing import Any, Callable, Optional, Union
 
-from ai_edge_torch import fx_pass_base
+from ai_edge_torch import fx_infra
 from ai_edge_torch.hlfb.mark_pattern import fx_utils
 import torch
 
@@ -118,7 +118,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 = fx_infra.run_passes(ep, [fx_infra.CanonicalizePass()])
       ep = ep.run_decompositions(decomp_table)
 
     scalar_locs = set()
@@ -152,13 +152,15 @@ class Pattern:
       self,
       name: str,
       module: Union[Callable, torch.nn.Module],
-      export_args: tuple[Any],
+      export_args: tuple[Any, ...],
       *,
       attr_builder: Callable[
           ["Pattern", GraphModule, InternalMatch], Optional[dict[str, Any]]
       ] = None,
       scalar_attr_trackers: list[ScalarAttrTracker] = None,
-      decomp_table: Optional[dict[torch._ops.OperatorBase, Callable]] = None,
+      extra_decomp_table: Optional[
+          dict[torch._ops.OperatorBase, Callable]
+      ] = None,
   ):
     """The PyTorch computation pattern to match against a model.
 
@@ -177,8 +179,9 @@ class Pattern:
       scalar_attr_trackers (list[ScalarAttrTracker]): the trackers 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.
+      extra_decomp_table (Optional[dict[torch._ops.OperatorBase, Callable]]):
+        Extra decomposition to be run on the pattern's exported program (in
+        addition to the default pre_convert_decomp in fx_infra.decomp).
     """
     if not isinstance(module, torch.nn.Module):
 
@@ -200,11 +203,14 @@ 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)
+
+    decomp_table = fx_infra.decomp.pre_convert_decomp()
+    if extra_decomp_table is not None:
+      decomp_table.update(extra_decomp_table)
+
+    exported_program = fx_infra.safe_run_decompositions(
+        exported_program, decomp_table
+    )
 
     self.exported_program = exported_program
     self.graph_module = self.exported_program.graph_module
@@ -222,6 +228,9 @@ class Pattern:
     # sanitization.
     self.graph_module = fx_utils.remove_clone_ops(self.graph_module)
     self.graph_module = fx_utils.remove_dangling_args(self.graph_module)
+    self.graph_module = fx_utils.remove_assert_tensor_metadata_nodes(
+        self.graph_module
+    )
 
     # Builds list of ordered input and output nodes.
     self.graph_nodes_map = {}

ai_edge_torch/hlfb/test/test_mark_pattern.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Tests for mark_pattern."""
 
+from ai_edge_torch import fx_infra
 from ai_edge_torch import lowertools
 from ai_edge_torch.hlfb import mark_pattern
 from ai_edge_torch.hlfb.mark_pattern import pattern as pattern_module
@@ -22,11 +23,13 @@ import torch
 from absl.testing import absltest as googletest
 
 
-def _export_stablehlo_mlir(model, args=None):
-  if not isinstance(model, torch.export.ExportedProgram):
-    ep = torch.export.export(model, args)
-  else:
-    ep = model
+def _export_and_decomp(mod, args):
+  ep = torch.export.export(mod, args)
+  ep = fx_infra.safe_run_decompositions(ep, fx_infra.decomp.pre_lower_decomp())
+  return ep
+
+
+def _to_mlir(ep: torch.export.ExportedProgram):
   return lowertools.exported_program_to_mlir_text(ep)
 
 
@@ -47,9 +50,9 @@ class TestMarkPattern(googletest.TestCase):
 
     model = TestModel().eval()
     args = (torch.rand(20, 20),)
-    exported_program = torch.export.export(model, args)
+    exported_program = _export_and_decomp(model, args)
     mark_pattern.mark_pattern(exported_program.graph_module, pattern)
-    mlir = _export_stablehlo_mlir(exported_program)
+    mlir = _to_mlir(exported_program)
 
     lowertools.assert_string_count(
         self,
@@ -73,9 +76,9 @@ class TestMarkPattern(googletest.TestCase):
 
     model = TestModel().eval()
     args = (torch.rand(20, 20),)
-    exported_program = torch.export.export(model, args)
+    exported_program = _export_and_decomp(model, args)
     mark_pattern.mark_pattern(exported_program.graph_module, pattern)
-    mlir = _export_stablehlo_mlir(exported_program)
+    mlir = _to_mlir(exported_program)
 
     lowertools.assert_string_count(
         self,
@@ -99,9 +102,9 @@ class TestMarkPattern(googletest.TestCase):
 
     model = TestModel().eval()
     args = (torch.rand(20, 20),)
-    exported_program = torch.export.export(model, args)
+    exported_program = _export_and_decomp(model, args)
     mark_pattern.mark_pattern(exported_program.graph_module, pattern)
-    mlir = _export_stablehlo_mlir(exported_program)
+    mlir = _to_mlir(exported_program)
 
     lowertools.assert_string_count(
         self,
@@ -137,9 +140,9 @@ class TestMarkPattern(googletest.TestCase):
 
     model = TestModel().eval()
     args = (torch.rand(10, 10),)
-    exported_program = torch.export.export(model, args)
+    exported_program = _export_and_decomp(model, args)
     mark_pattern.mark_pattern(exported_program.graph_module, pattern)
-    mlir = _export_stablehlo_mlir(exported_program)
+    mlir = _to_mlir(exported_program)
 
     lowertools.assert_string_count(
         self,
@@ -169,9 +172,9 @@ class TestMarkPattern(googletest.TestCase):
 
     model = TestModel().eval()
     args = (torch.rand(20, 20), torch.rand(20, 20))
-    exported_program = torch.export.export(model, args)
+    exported_program = _export_and_decomp(model, args)
     mark_pattern.mark_pattern(exported_program.graph_module, pattern)
-    mlir = _export_stablehlo_mlir(exported_program)
+    mlir = _to_mlir(exported_program)
 
     lowertools.assert_string_count(
         self,

ai_edge_torch/odml_torch/_torch_future.py

@@ -59,30 +59,3 @@ 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)
-
-
-def dummy_decomp_table():
-  """Build dummy decomp table for run_decompositions without any decompositions.
-
-  Compatible for torch<=2.5.
-
-  Returns:
-    Decomp table for ExportedProgram.run_decompositions.
-  """
-  return {
-      torch._ops.OperatorBase(): lambda: None,
-  }

ai_edge_torch/odml_torch/export.py

@@ -20,6 +20,7 @@ import io
 import operator
 from typing import Any, Callable, Optional
 
+from ai_edge_torch import fx_infra
 from jax.lib import xla_extension
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import func
@@ -302,16 +303,13 @@ def exported_program_to_mlir(
     exported_program: torch.export.ExportedProgram,
 ) -> MlirLowered:
   """Lower the exported program to MLIR."""
-  exported_program = _torch_future.safe_run_decompositions(
-      exported_program, lowerings.decompositions()
+  exported_program = fx_infra.safe_run_decompositions(
+      exported_program,
+      fx_infra.decomp.pre_lower_decomp(),
   )
-
   _convert_i64_to_i32(exported_program)
-
-  # No decompositions but just retracing/cananicalization.
-  exported_program = _torch_future.safe_run_decompositions(
-      exported_program, _torch_future.dummy_decomp_table()
-  )
+  # Run decompositions for retracing and cananicalization.
+  exported_program = fx_infra.safe_run_decompositions(exported_program, {})
 
   # Passes below mutate the exported program to a state not executable by torch.
   # Do not call run_decompositions after applying the passes.

ai_edge_torch/odml_torch/lowerings/__init__.py

@@ -15,6 +15,7 @@
 from . import _basic
 from . import _batch_norm
 from . import _convolution
+from . import _decomp_registry
 from . import _jax_lowerings
 from . import _layer_norm
 from . import _quantized_decomposed
@@ -22,6 +23,5 @@ from . import _rand
 from . import context
 from . import registry
 from . import utils
-from .decomp import decompositions
 from .registry import lookup
 from .registry import lower

ai_edge_torch/odml_torch/lowerings/_decomp_registry.py

@@ -0,0 +1,65 @@
+# 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.
+# ==============================================================================
+"""Torch export decompositions to run before lowering."""
+
+from ai_edge_torch import fx_infra
+import torch
+
+
+fx_infra.decomp.update_pre_lower_decomp(
+    torch._decomp.get_decompositions([
+        torch.ops.aten.upsample_nearest2d,
+        torch.ops.aten._native_batch_norm_legit.no_stats,
+        torch.ops.aten._native_batch_norm_legit_functional,
+        torch.ops.aten._adaptive_avg_pool2d,
+        torch.ops.aten._adaptive_avg_pool3d,
+        torch.ops.aten.grid_sampler_2d,
+        torch.ops.aten.native_group_norm,
+        torch.ops.aten.native_dropout,
+        torch.ops.aten.reflection_pad1d,
+        torch.ops.aten.reflection_pad2d,
+        torch.ops.aten.reflection_pad3d,
+        torch.ops.aten.replication_pad1d,
+        torch.ops.aten.replication_pad2d,
+        torch.ops.aten.replication_pad3d,
+        torch.ops.aten.addmm,
+    ])
+)
+
+fx_infra.decomp.remove_pre_lower_decomp(torch.ops.aten.roll)
+
+# Torch's default einsum impl/decompositions is less efficient and
+# optimized through converter than JAX's impl. Disable einsum
+# decomposition to use JAX bridge for a more efficient lowering.
+fx_infra.decomp.remove_pre_lower_decomp(torch.ops.aten.einsum.default)
+
+
+# Override noop aten op decompositions for faster run_decompositions.
+fx_infra.decomp.add_pre_convert_decomp(
+    torch.ops.aten.alias.default, lambda x: x
+)
+fx_infra.decomp.add_pre_convert_decomp(
+    torch.ops.aten.detach.default, lambda x: x
+)
+
+# Override _safe_softmax decompositions with regular softmax.
+# _safe_softmax introduces additional check-select ops to guard extreme
+# input values to softmax, which could make the converted model inefficient
+# on-device.
+if hasattr(torch.ops.aten, "_safe_softmax"):
+  fx_infra.decomp.add_pre_convert_decomp(
+      torch.ops.aten._safe_softmax.default,
+      torch.softmax,
+  )

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20250121.dist-info → ai_edge_torch_nightly-0.3.0.dev20250123.dist-info}/METADATA

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

@@ -1,32 +1,32 @@
-ai_edge_torch/__init__.py,sha256=rq9ZtMJLG8yYNC4tNE4rpl94UAUClZW7f4GAr6HBVDQ,1208
+ai_edge_torch/__init__.py,sha256=8sPR_5uXJA4NEE0nIwNdSl-ADOJEoR8hAgYvBQDY70Y,1208
 ai_edge_torch/_config.py,sha256=PKtOtBOup-cM0wBdQxby6HzuhLhIC3oq-TBG8FF4znE,2161
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
-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=Zlgc-dAVs0gcH1g86CJVS-yy-T9eJBuCXDFHRUscOVA,706
+ai_edge_torch/version.py,sha256=szrxg2aB7mcm59IL_QVIqapmbw9Nz8AQ28vc9684bqY,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/_convert/conversion.py,sha256=pSDY0CzZQP_jAMjSfQ1O7Ud_AF5ZDeDF-nE3nAu_hoo,5815
+ai_edge_torch/_convert/conversion.py,sha256=mckvxznKLXdF2HuJg_IxQaT5Ty-iWl_iXElHEugH3VI,5452
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
 ai_edge_torch/_convert/converter.py,sha256=075F8LRewk_033Ebsnft7FJr3KgtIbtZ_-8udIPy6ho,9980
 ai_edge_torch/_convert/signature.py,sha256=-YKJdLk-eNEHfhdPCtcQVtZf915SoVePEFxKXPPf16c,2572
 ai_edge_torch/_convert/to_channel_last_io.py,sha256=_31phf7TYgZY2ftpNbrdlB1RhDium1lz_BXEQ6IsMFc,2893
-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=qb4JBDi4Xca14JJUIcaaZQIJiyqKyHJF49jsRCIFCVA,4335
-ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=NxT-iCOHq3r3jeZ8qhNoPXV5w8l2eRMu4yEcBri3NxY,2398
-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/__init__.py,sha256=dG4WIICk0FqCH9euvbYHHsybRN7B1cYcuxN_OYxmjWo,1263
+ai_edge_torch/_convert/fx_passes/build_aten_composite_pass.py,sha256=a1KhqLetFb_efRHjX4T-zH0vF-U37Ha5I1CPIAsIluE,9211
+ai_edge_torch/_convert/fx_passes/build_interpolate_composite_pass.py,sha256=3JyjiHpn17Zhfq3yGQXK5LMH71DQPXHb_4GOkP9uAjY,4251
+ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=Z6E3U7SYZvMl3Ivpqa3burVOLKFndEZuNmWKNxjq2mM,2386
+ai_edge_torch/_convert/fx_passes/remove_non_user_outputs_pass.py,sha256=HCOkj0k3NhaYbtfjE8HDXVmYhZ9fL5V_u6VunVh9mN4,2116
+ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/__init__.py,sha256=UKC-wM93-oe8spxyFqgybJ0TwnSRw8f-SOA2glCh2FA,890
+ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/_decomp_registry.py,sha256=aWO_zHDF4j_hokoKJQNFIFmua4ysXztsgS6pcyBUht0,1082
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_check.py,sha256=S_Bniv6jY16oOoFUzlyECQ0I2HDjG2D1MOI-QYPk3jQ,8061
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_mark.py,sha256=4RyGUwR22bZqkn_TnptenFJodc_Q43f4_SBG7gmTbos,1621
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py,sha256=zoAZ2TXKvxUnWnT11U4tx2uF0J5kkNXydgaW7JzfkXI,13811
-ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/op_func_registry.py,sha256=bsYnudRlXp1PJlu4GF25KSogSkBGQPSaecBrUTONKaw,1031
-ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/pass_body.py,sha256=t94Am3iPbYQekg-rrtc-jS_aDWtEgAAj7pAKHrG0-9U,10563
+ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/op_func_registry.py,sha256=OhisegHY2j4cv_m9auCh9Mq9qmm1lUqpFLVO9X-oBlc,1032
+ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/pass_body.py,sha256=lgoH32l6zAbWTCpa_4-RWkHjqbNaPsBnhSObLIX8dL4,10551
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/utils.py,sha256=YLMttMg5PdvXTtQ8lxpKb434UGVvYVALV1-xeuH4UGc,2131
 ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/__init__.py,sha256=D8VX8SbCzfoyvPgMFHK7yxD7R-bzLxp2gfdKxgrWekA,742
 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=gK9QJuLbpjXt0l6tVnzl9Miq6GLkJR-hB67i3VE13Og,17224
+ai_edge_torch/_convert/test/test_convert.py,sha256=o6tuJkD-ESaQxLxJpN104qpchm3LCtPmHinzQxe6PSg,17226
 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
@@ -37,6 +37,12 @@ ai_edge_torch/debug/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf
 ai_edge_torch/debug/test/test_culprit.py,sha256=fRN-8jJicawJ2mhPRQNAQUZ8AdGg-s0tYMXyhnLAlWw,3875
 ai_edge_torch/debug/test/test_search_model.py,sha256=-RuU0QsjqkfzZF2IbeA55MoeVOawhbgiSEu96PmioPE,1668
 ai_edge_torch/experimental/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
+ai_edge_torch/fx_infra/__init__.py,sha256=APjkSqEfwDxcnI8k53rGi3Ef-G2L-M8fdaPGpxXtuiI,1347
+ai_edge_torch/fx_infra/_canonicalize_pass.py,sha256=GDRoDdPVQw--QQFTT5J_C3TVuphL31m6K6F1-67SE4s,1097
+ai_edge_torch/fx_infra/_safe_run_decompositions.py,sha256=ZbWheeZ8ydsxCk2aVGUgUynrkEkBOMjBCzPhS5uq4sU,2595
+ai_edge_torch/fx_infra/decomp.py,sha256=S58SCgwMHYVFl_hJwlJxvu2wcI-AGNn82gel3qmTPrU,2500
+ai_edge_torch/fx_infra/graph_utils.py,sha256=3UZAOHWOUh2LCj1E2_AKQn3gRDILi9JCdqSScjyOd4M,1535
+ai_edge_torch/fx_infra/pass_base.py,sha256=Ic2AlhSoRFscz6l7gJKvWVNMDLQFfAw5kRf84-ZR9qM,2904
 ai_edge_torch/generative/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/amd_llama_135m/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -115,8 +121,8 @@ ai_edge_torch/generative/examples/tiny_llama/__init__.py,sha256=hHLluseD2R0Hh4W6
 ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=VU0c5pgvrUtaTboT1xuDBGjpKOM85aqtaB_hYfSBuEk,2544
 ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=mhJ18rb9sxrYRzv1YSzhbNs97oUZck99avZDcUO2oV8,2800
 ai_edge_torch/generative/examples/tiny_llama/verify.py,sha256=7Bk8z033M-BCXJ299fpQNXYAudBbZoDQp9934xcvg50,2426
-ai_edge_torch/generative/fx_passes/__init__.py,sha256=jrzCB3ZyY_t5jJM1e2Czdt3DjAIL43R0_a-T-I7wOzw,1155
-ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=hhxSQvkDMv0isZJhmuLiod66ZODaJ8uSPSVTJVHBabQ,1931
+ai_edge_torch/generative/fx_passes/__init__.py,sha256=4rFrppMRKlTwwZeX1ON_cdp4yUqoTOES161IZQkJF6c,1143
+ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=n5TbXdhBZi8jQe4j7-rox_MugMVvW8ReOhkTA3pfQkw,1919
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/attention.py,sha256=GrAy8CT1pEsgRoB8JQP6PlnNYk8kQ4U3YANfSiTJKn8,13776
 ai_edge_torch/generative/layers/attention_utils.py,sha256=zBVwlBUTs-nStIKCZG0ks5ra7tsqc9ShfakFJKH5rds,7344
@@ -159,11 +165,11 @@ ai_edge_torch/generative/utilities/t5_loader.py,sha256=tEsfy8-ymzbbjOIc-oesXF3yG
 ai_edge_torch/generative/utilities/transformers_verifier.py,sha256=8sp9m_FMcXn7nqOrochtu2jIANkJKhnhIBUmH0ZTDR4,1549
 ai_edge_torch/generative/utilities/verifier.py,sha256=6lnBU9Cy5GanB8JWK3-2_VU3PxqunDWGe-SgSLba5Yw,12065
 ai_edge_torch/hlfb/__init__.py,sha256=sH4um75na-O8tzxN6chFyp6Y4xnexsE7kUQpZySv6dE,735
-ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=-BYE7MGMxr-VfBy8tAiiOaCqYv8ytJ0w5l2P8B7h3eM,5387
-ai_edge_torch/hlfb/mark_pattern/fx_utils.py,sha256=taWLpF5IVglxlsF9HM2dIoKDXuQREaCRAXtJeG5gKzs,2073
-ai_edge_torch/hlfb/mark_pattern/pattern.py,sha256=7bv9XqRkm1pjxiVL4Cm1cArExnolId8hQKFHtvlkCI8,10061
+ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=JsVmYrM_JEuN_smMHXUsRlo3Liapp7UyktbPpPARwDk,5386
+ai_edge_torch/hlfb/mark_pattern/fx_utils.py,sha256=YCtMgu-4w2BQ5fpnlpWC6IauKPf_tVqc7Ff91OTqlSw,1796
+ai_edge_torch/hlfb/mark_pattern/pattern.py,sha256=Ui6BrehF3zJJN7uTxKwbO2yCY9mYjbewlQzAxzZv9Es,10274
 ai_edge_torch/hlfb/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/hlfb/test/test_mark_pattern.py,sha256=-5UqJyk__1YbUNGuxi4b2sn0CED0W-G337AXwxPGdEs,5567
+ai_edge_torch/hlfb/test/test_mark_pattern.py,sha256=5kmOJWCc7sU1Hrqr1y17BtShUrssTfaV1sMyUvdMbsg,5573
 ai_edge_torch/lowertools/__init__.py,sha256=A8WBXvWtuFYYWtNTqPD7waVntLaSVAnSMwx5ugjZBIw,761
 ai_edge_torch/lowertools/_shim.py,sha256=Mbg16tnCVK0YsHowfbpqpNX1qySuMLvpGI_-I5SIrG0,3276
 ai_edge_torch/lowertools/common_utils.py,sha256=4HQtquPZ6oiId8vR_1ykW_uK4ELnyo5zo3MlX1QYW4c,4513
@@ -172,9 +178,9 @@ ai_edge_torch/lowertools/test_utils.py,sha256=mdxTlhqHABZEQ_GEmPFCL8LIAWtqRtYZUG
 ai_edge_torch/lowertools/torch_xla_utils.py,sha256=1EytIw2R6dthhLhf69wN1L9BaQTeybCD0wga-PhHcMI,9518
 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=6fMr6C6rP3roVRFQxBH9Yr7656WtodZuNvhxLzKad_g,3320
+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=xaJP_dca7P1sC0KvAI5ExxklIoW9nCrRM3vsdvhVfm8,13451
+ai_edge_torch/odml_torch/export.py,sha256=YN7QPrQ8W6T3YVOdyIGadfSQuBroMjIqAMB9FeUa7Ho,13447
 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
@@ -186,16 +192,16 @@ ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py,sha256=IvOBQyROI9WHS3umHRxsDW
 ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=e9Oa4J3An9FYr3zM0OzjzyNNitEeg-IoBUGNfUxsDSA,798
 ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=LqwZ1vCJTSOzgzvH8LUAN-sAkF-l_pGj1AMEIzAqHCA,6638
 ai_edge_torch/odml_torch/jax_bridge/utils.py,sha256=T8isGc896VrHZ6c_L5pYmLpolQ7ibcOlgWfPuVFPzIg,2264
-ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=GWFl7WWgExLXu6FEYxnig5_g6hd_Sfnl8690uFg2-CU,1013
+ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=uJ-niilt1c-D6QJzLwgvCUf62le_JsxQTlqj_iP_Ps0,1009
 ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=8mZTp_ybcMO3tDRQdlDP68BVeTw560XsTR4XH-ldTdc,9987
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=Q0aDzyUcZMoSzSbOU-r3LJMgPe6fble0QwdYVIOHHHk,6887
+ai_edge_torch/odml_torch/lowerings/_decomp_registry.py,sha256=VhmeGFnB5hrUsALiVWV96JJOqPDrTIWouHjTvLuT5eU,2477
 ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=CJHWkmY4aAVQ5dmFsVc3Ox9TPkoLSNOfa96psD4CLRo,11561
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=khJIvDVk2s332Nd2Be-5dM6-wp5DGff61HCV5lskHmQ,3011
 ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=XDZ0zLej_XaQDJnaAAxhNFAd7NfQm5SOVEp_nno_krA,6178
 ai_edge_torch/odml_torch/lowerings/_rand.py,sha256=g6SuqDkuC6hD35lyP1-5H7ASDIzPSmKukeNT5naZSv8,4133
 ai_edge_torch/odml_torch/lowerings/context.py,sha256=jslcCv7r_HtImSRTxJwHAUV_QCu9Jub51lovmoBkmFA,1295
-ai_edge_torch/odml_torch/lowerings/decomp.py,sha256=0dIkVNN9ubgfwSxeaBFfTmWarWWG8Q1M0HX_1FShHik,2610
 ai_edge_torch/odml_torch/lowerings/registry.py,sha256=Tp2h11l5uTifO0aIkuUOWAF_ibEjmd65Xx99w3EXuGE,1924
 ai_edge_torch/odml_torch/lowerings/utils.py,sha256=pqM6mumpviFDHRaabp93CUAngzEZmWcAHl0nTDgyI2g,6167
 ai_edge_torch/odml_torch/passes/__init__.py,sha256=AVwIwUTMx7rXacKjGy4kwrtMd3XB2v_ncdc40KOjUqQ,1245
@@ -206,8 +212,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.dev20250121.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20250121.dist-info/METADATA,sha256=NWhBDavlpQ2NnaCILfBYuSK6xCLz9YOAFzDH7Ydrlgg,1966
-ai_edge_torch_nightly-0.3.0.dev20250121.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20250121.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20250121.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20250123.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20250123.dist-info/METADATA,sha256=1IZCBOcKVCWbEfAQvEMgt39cuATDIzpK6AhW_gTnIY4,1966
+ai_edge_torch_nightly-0.3.0.dev20250123.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20250123.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20250123.dist-info/RECORD,,

ai_edge_torch/odml_torch/lowerings/decomp.py

@@ -1,69 +0,0 @@
-# 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.
-# ==============================================================================
-"""Torch export decompositions to run before lowering."""
-
-import functools
-
-import torch
-
-
-@functools.cache
-def decompositions():
-  # Base: Core ATen decompositions
-  decompositions = torch._decomp.core_aten_decompositions()
-
-  decompositions.update(
-      torch._decomp.get_decompositions([
-          torch.ops.aten.upsample_nearest2d,
-          torch.ops.aten._native_batch_norm_legit.no_stats,
-          torch.ops.aten._native_batch_norm_legit_functional,
-          torch.ops.aten._adaptive_avg_pool2d,
-          torch.ops.aten._adaptive_avg_pool3d,
-          torch.ops.aten.grid_sampler_2d,
-          torch.ops.aten.native_group_norm,
-          torch.ops.aten.native_dropout,
-          torch.ops.aten.reflection_pad1d,
-          torch.ops.aten.reflection_pad2d,
-          torch.ops.aten.reflection_pad3d,
-          torch.ops.aten.replication_pad1d,
-          torch.ops.aten.replication_pad2d,
-          torch.ops.aten.replication_pad3d,
-          torch.ops.aten.addmm,
-      ])
-  )
-
-  torch._decomp.remove_decompositions(
-      decompositions,
-      [
-          torch.ops.aten.roll,
-          # Torch's default einsum impl/decompositions is less efficient and
-          # optimized through converter than JAX's impl. Disable einsum
-          # decomposition to use JAX bridge for a more efficient lowering.
-          torch.ops.aten.einsum.default,
-      ],
-  )
-
-  # Override noop aten op decompositions for faster run_decompositions.
-  decompositions[torch.ops.aten.alias.default] = lambda x: x
-  decompositions[torch.ops.aten.detach.default] = lambda x: x
-
-  # Override _safe_softmax decompositions with regular softmax.
-  # _safe_softmax introduces additional check-select ops to guard extreme
-  # input values to softmax, which could make the converted model inefficient
-  # on-device.
-  if hasattr(torch.ops.aten, "_safe_softmax"):
-    decompositions[torch.ops.aten._safe_softmax.default] = torch.softmax
-
-  return decompositions