ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl

ai_edge_torch/convert/converter.py

@@ -0,0 +1,202 @@
+# 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 __future__ import annotations
+
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+
+from ai_edge_torch import model
+from ai_edge_torch.convert import conversion
+from ai_edge_torch.convert import conversion_utils as cutils
+from ai_edge_torch.quantize import quant_config as qcfg
+
+
+class Converter:
+
+  def __init__(self):
+    self._signatures: list[cutils.Signature] = []
+
+  def signature(
+      self,
+      name: str,
+      module: torch.nn.Module,
+      sample_args=None,
+      sample_kwargs=None,
+      *,
+      dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+  ) -> Converter:
+    """Alias to `add_signature`"""
+    return self.add_signature(
+        name, module, sample_args, sample_kwargs, dynamic_shapes=dynamic_shapes
+    )
+
+  def add_signature(
+      self,
+      name: str,
+      module: torch.nn.Module,
+      sample_args=None,
+      sample_kwargs=None,
+      *,
+      dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+  ) -> Converter:
+    """Allows adding a new named torch model along with sample args to the conversion.
+
+    Args:
+      name: The name of the signature included in the converted edge model.
+      module: The torch module to be converted.
+      sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
+      sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to conversion.
+      dynamic_shapes: Optional dict or tuple that specify dynamic shape specifications for each input in original order.
+        See https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
+
+    Raises:
+      ValueError: If a signature with the provided name already exists.
+    """
+
+    if name in [sig.name for sig in self._signatures]:
+      raise ValueError(f"A signature with the provided name ({name}) is already added.")
+
+    if sample_args is None and sample_kwargs is None:
+      raise ValueError("sample_args or sample_kwargs must be provided.")
+
+    self._signatures.append(
+        cutils.Signature(
+            name, module, sample_args, sample_kwargs, dynamic_shapes=dynamic_shapes
+        )
+    )
+    return self
+
+  def convert(
+      self,
+      module: torch.nn.Module = None,
+      sample_args=None,
+      sample_kwargs=None,
+      *,
+      quant_config: Optional[qcfg.QuantConfig] = None,
+      dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+      _ai_edge_converter_flags: dict = {},
+  ) -> model.TfLiteModel:
+    """Finalizes the conversion and produces an edge model.
+
+    This could be called with no arguments as follows:
+
+      edge_model = Converter().signature(name, module, args).convert()
+
+    Or it could be used to set the default signature for the converted edge model:
+
+      edge_model =  Converter().convert(module, args)
+
+    Args:
+      name: The name of the signature included in the converted edge model.
+      module: The torch module to be converted.
+      sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
+      sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to 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 https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
+      _ai_edge_converter_flags: A nested dictionary allowing setting flags for the underlying converter.
+        This gives access to an implementation detail of this function and so needs to be treated as such.
+        Please do not rely on this parameter except for local debugging as this can be removed in a future release.
+
+    Raises:
+      ValueError: If the arguments are not provided as expected. See the example in this functions's comment.
+    """
+    if module is not None:
+      if (
+          sample_args is not None or sample_kwargs is not None
+      ):  # both module and args provided
+        self.add_signature(
+            cutils.DEFAULT_SIGNATURE_NAME,
+            module,
+            sample_args,
+            sample_kwargs,
+            dynamic_shapes=dynamic_shapes,
+        )
+      else:  # module is provided but not args
+        raise ValueError(
+            "sample_args or sample_kwargs must be provided if a module is specified."
+        )
+
+    return conversion.convert_signatures(
+        self._signatures,
+        quant_config=quant_config,
+        _tfl_converter_flags=_ai_edge_converter_flags,
+    )
+
+
+def signature(
+    name: str,
+    module: torch.nn.Module,
+    sample_args=None,
+    sample_kwargs=None,
+    dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+) -> Converter:
+  """Initiates a Converter object with the provided signature.
+
+  Args:
+    name: The name of the signature included in the converted edge model.
+    module: The torch module to be converted.
+    sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
+    sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to conversion.
+    dynamic_shapes: Optional dict or tuple that specify dynamic shape specifications for each input in original order.
+      See https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
+
+  Example:
+    converter = ai_edge_torch.signature(name, module, args)
+    edge_model = converter.convert()
+
+  """
+  return Converter().signature(
+      name, module, sample_args, sample_kwargs, dynamic_shapes=dynamic_shapes
+  )
+
+
+def convert(
+    module: torch.nn.Module = None,
+    sample_args=None,
+    sample_kwargs=None,
+    *,
+    quant_config: Optional[qcfg.QuantConfig] = None,
+    dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+    _ai_edge_converter_flags: dict = {},
+) -> model.TfLiteModel:
+  """Allows converting a PyTorch model to an edge model with one default signature in one step.
+
+  Args:
+    module: The torch module to be converted.
+    sample_args: Tuple of tensors by which the torch module will be traced with prior to conversion.
+    sample_kwargs: Dict of str to tensor by which the torch module will be traced with prior to 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 https://pytorch.org/docs/stable/export.html#expressing-dynamism for more details.
+    _ai_edge_converter_flags: A nested dictionary allowing setting flags for the underlying converter.
+      This gives access to an implementation detail of this function and so needs to be treated as such.
+      Please do not rely on this parameter except for local debugging as this can be removed in a future release.
+
+  Example:
+    edge_model = ai_edge_torch.convert(module, args)
+
+  """
+
+  return Converter().convert(
+      module,
+      sample_args,
+      sample_kwargs,
+      quant_config=quant_config,
+      dynamic_shapes=dynamic_shapes,
+      _ai_edge_converter_flags=_ai_edge_converter_flags,
+  )

ai_edge_torch/convert/fx_passes/__init__.py

@@ -0,0 +1,59 @@
+# 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 Sequence, Union
+
+from torch.export import ExportedProgram
+from torch.fx.passes.infra.pass_manager import pass_result_wrapper
+import torch.utils._pytree as pytree
+
+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.convert.fx_passes._pass_base import FxPassBase
+from ai_edge_torch.convert.fx_passes._pass_base import FxPassResult
+from ai_edge_torch.convert.fx_passes.build_aten_composite_pass import BuildAtenCompositePass  # NOQA
+from ai_edge_torch.convert.fx_passes.build_interpolate_composite_pass import BuildInterpolateCompositePass  # NOQA
+from ai_edge_torch.convert.fx_passes.canonicalize_pass import CanonicalizePass
+from ai_edge_torch.convert.fx_passes.inject_mlir_debuginfo_pass import InjectMlirDebuginfoPass  # NOQA
+from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import OptimizeLayoutTransposesPass  # NOQA
+
+
+# TODO(cnchan): make a PassManager class.
+def run_passes(
+    exported_program: ExportedProgram,
+    passes: Sequence[Union[ExportedProgramPassBase, FxPassBase]],
+) -> ExportedProgram:
+  passes, _ = pytree.tree_flatten(passes)
+  for pass_ in passes:
+    if not isinstance(pass_, ExportedProgramPassBase):
+      pass_ = pass_result_wrapper(pass_)
+    if isinstance(pass_, ExportedProgramPassBase):
+      exported_program = pass_(exported_program).exported_program
+    else:
+      gm = exported_program.graph_module
+      gm, modified = pass_(gm)
+      if modified and gm is not exported_program.graph_module:
+        exported_program = ExportedProgram(
+            root=gm,
+            graph=gm.graph,
+            graph_signature=exported_program.graph_signature,
+            state_dict=exported_program.state_dict,
+            range_constraints=exported_program.range_constraints,
+            module_call_graph=exported_program.module_call_graph,
+            example_inputs=exported_program.example_inputs,
+            verifier=exported_program.verifier,
+            constants=exported_program.constants,
+        )
+  return exported_program

ai_edge_torch/convert/fx_passes/_pass_base.py

@@ -0,0 +1,49 @@
+# 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 abc
+from collections import namedtuple
+
+import torch
+from torch.export import ExportedProgram
+from torch.fx.passes.infra.pass_base import PassBase as FxPassBase
+from torch.fx.passes.infra.pass_base import PassResult as FxPassResult
+
+
+class ExportedProgramPassResult(
+    namedtuple("ExportedProgramPassResult", ["exported_program", "modified"])
+):
+
+  def __new__(cls, exported_program, modified):
+    return super().__new__(cls, exported_program, modified)
+
+
+class ExportedProgramPassBase(abc.ABC):
+
+  def __call__(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
+    self.requires(exported_program)
+    res = self.call(exported_program)
+    self.ensures(exported_program)
+    return res
+
+  @abc.abstractmethod
+  def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
+    pass
+
+  def requires(self, exported_program: ExportedProgram) -> None:
+    pass
+
+  def ensures(self, exported_program: ExportedProgram) -> None:
+    pass

ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py

@@ -0,0 +1,225 @@
+# 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 copy
+import functools
+from typing import Any, Callable
+
+import torch
+from torch.fx import GraphModule
+from torch.fx import Node
+from torch.fx.passes.infra.pass_base import PassBase
+from torch.fx.passes.infra.pass_base import PassResult
+import torch.utils._pytree as pytree
+
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
+
+_composite_builders: dict[Callable, Callable[[GraphModule, Node], None]] = {}
+
+
+def _register_composite_builder(op):
+  def inner(func):
+    if isinstance(op, torch._ops.OpOverloadPacket):
+      for overload in v.overloads():
+        _composite_builders[getattr(v, overload)] = func
+    else:
+      _composite_builders[op] = func
+    return func
+
+  return inner
+
+
+def _tree_map_to_composite_attr_values(values, *, stringify_incompatible_values=True):
+
+  def convert(value):
+    nonlocal stringify_incompatible_values
+    if value is None:
+      return "py_None"
+    if isinstance(value, (str, int, float, bool)):
+      return value
+
+    if stringify_incompatible_values:
+      return str(value)
+    return value
+
+  return pytree.tree_map(convert, values)
+
+
+class TorchOpArgumentsMapper:
+
+  def __init__(self, op):
+    if isinstance(op, torch._ops.OpOverloadPacket):
+      op = op.default
+
+    assert hasattr(op, "_schema")
+    self.op = op
+    self.arg_specs = [(spec.name, spec.default_value) for spec in op._schema.arguments]
+
+  def get_full_kwargs(self, args, kwargs=None) -> dict[str, Any]:
+    """Inspect the op's schema and extract all its args and kwargs
+    into one single kwargs dict, with default values for those
+    unspecified args and kwargs.
+    """
+    full_kwargs = {**(kwargs or {})}
+
+    for arg, (name, default_value) in zip(args, self.arg_specs):
+      full_kwargs[name] = arg
+
+    for name, default_value in self.arg_specs[len(args) :]:
+      if name not in full_kwargs:
+        full_kwargs[name] = default_value
+
+    return full_kwargs
+
+
+@_register_composite_builder(torch.ops.aten.hardswish.default)
+def _aten_hardswish(gm: GraphModule, node: Node):
+  op = node.target
+
+  def hardswish(self: torch.Tensor):
+    nonlocal op
+    builder = StableHLOCompositeBuilder("aten.hardswish.default")
+    self = builder.mark_inputs(self)
+    output = op(self)
+    output = builder.mark_outputs(output)
+    return output
+
+  node.target = hardswish
+
+
+@_register_composite_builder(torch.ops.aten.gelu.default)
+def _aten_gelu(gm: GraphModule, node: Node):
+  op = node.target
+  args_mapper = TorchOpArgumentsMapper(op)
+
+  def gelu(*args, **kwargs):
+    nonlocal op, args_mapper
+
+    full_kwargs = args_mapper.get_full_kwargs(args, kwargs)
+
+    # TFLite supports exact and tanh approximate.
+    if full_kwargs["approximate"] != "none" and full_kwargs["approximate"] != "tanh":
+      return op(*args, **kwargs)
+
+    builder = StableHLOCompositeBuilder(
+        "aten.gelu.default",
+        attr=_tree_map_to_composite_attr_values(
+            {
+                "approximate": full_kwargs["approximate"],
+            }
+        ),
+    )
+    full_kwargs["self"] = builder.mark_inputs(full_kwargs["self"])
+    output = op(full_kwargs["self"])
+    output = builder.mark_outputs(output)
+    return output
+
+  node.target = gelu
+
+
+@_register_composite_builder(torch.ops.aten.avg_pool2d.default)
+def _aten_avg_pool2d(gm: GraphModule, node: Node):
+  op = node.target
+  args_mapper = TorchOpArgumentsMapper(op)
+
+  def avg_pool2d(*args, **kwargs):
+    nonlocal op, args_mapper
+
+    full_kwargs = args_mapper.get_full_kwargs(args, kwargs)
+
+    def is_same_padding(
+        input_shape: list[int],
+        kernel_size: list[int],
+        stride: list[int],
+        padding: list[int],
+    ):
+      for dim_input_size, dim_kernel_size, dim_stride, dim_padding in zip(
+          input_shape, kernel_size, stride, padding
+      ):
+        dim_output_size = int((dim_input_size + dim_stride - 1) / dim_stride)
+        padding_needed = max(
+            0, (dim_output_size - 1) * dim_stride + dim_kernel_size - dim_input_size
+        )
+        if padding_needed % 2 != 0:
+          return False
+
+        if padding_needed // 2 != dim_padding:
+          return False
+      return True
+
+    def is_valid_padding(padding: list[int]):
+      return not any(padding)
+
+    # We prefer to avoid passing empty arrays to composite attributes
+    # as they will be lowered to an ArrayAttr so canonicalizing according
+    # to the default behaviour here.
+    if not full_kwargs["stride"]:
+      full_kwargs["stride"] = full_kwargs["kernel_size"]
+
+    # Only wrap in a composite when the underlying converter can handle it.
+    # TODO We should be able to remove this if the converter can inline composites when it can not handle them.
+
+    # We don't cover any cases where the divisor_override is set.
+    if full_kwargs["divisor_override"] is not None:
+      return op(*args, **kwargs)
+
+    if full_kwargs["ceil_mode"] and not full_kwargs["count_include_pad"]:
+      return op(*args, **kwargs)
+
+    # We also can not cover a case where count_include_pad is False but the padding is custom.
+    if (
+        not full_kwargs["count_include_pad"]
+        and not is_valid_padding(full_kwargs["padding"])
+        and not is_same_padding(
+            list(full_kwargs["self"].shape)[2:],
+            full_kwargs["kernel_size"],
+            full_kwargs["stride"],
+            full_kwargs["padding"],
+        )
+    ):
+      return op(*args, **kwargs)
+
+    builder = StableHLOCompositeBuilder(
+        "aten.avg_pool2d.default",
+        attr=_tree_map_to_composite_attr_values(
+            {
+                "kernel_size": full_kwargs["kernel_size"],
+                "stride": full_kwargs["stride"],
+                "padding": full_kwargs["padding"],
+                "ceil_mode": full_kwargs["ceil_mode"],
+                "count_include_pad": full_kwargs["count_include_pad"],
+                "divisor_override": full_kwargs["divisor_override"],
+            }
+        ),
+    )
+
+    full_kwargs["self"] = builder.mark_inputs(full_kwargs["self"])
+    output = op(**full_kwargs)
+    output = builder.mark_outputs(output)
+    return output
+
+  node.target = avg_pool2d
+
+
+class BuildAtenCompositePass(PassBase):
+
+  def call(self, graph_module: GraphModule):
+    for node in graph_module.graph.nodes:
+      if node.target in _composite_builders:
+        _composite_builders[node.target](graph_module, node)
+
+    graph_module.graph.lint()
+    graph_module.recompile()
+    return PassResult(graph_module, True)

ai_edge_torch/convert/fx_passes/build_interpolate_composite_pass.py

@@ -0,0 +1,123 @@
+# 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 functools
+
+import torch
+
+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():
+  pattern = mark_pattern.Pattern(
+      "odml.upsample_bilinear2d",
+      lambda x: torch.nn.functional.interpolate(
+          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
+  def attr_builder(pattern, graph_module, internal_match):
+    output = internal_match.returning_nodes[0]
+    output_h, output_w = output.meta["val"].shape[-2:]
+    return {
+        "output": (int(output_h), int(output_w)),
+        "align_corners": False,
+    }
+
+  return pattern
+
+
+@functools.cache
+def _get_upsample_bilinear2d_align_corners_pattern():
+  pattern = mark_pattern.Pattern(
+      "odml.upsample_bilinear2d",
+      lambda x: torch.nn.functional.interpolate(
+          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
+  def attr_builder(graph_module, pattern, internal_match):
+    output = internal_match.returning_nodes[0]
+    output_h, output_w = output.meta["val"].shape[-2:]
+    return {
+        "output": (int(output_h), int(output_w)),
+        "align_corners": True,
+    }
+
+  return pattern
+
+
+@functools.cache
+def _get_interpolate_nearest2d_pattern():
+  pattern = mark_pattern.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
+  def attr_builder(pattern, graph_module, internal_match):
+    output = internal_match.returning_nodes[0]
+    output_h, output_w = output.meta["val"].shape[-2:]
+    return {
+        "size": (int(output_h), int(output_w)),
+        "is_nchw_op": True,
+    }
+
+  return pattern
+
+
+class BuildInterpolateCompositePass(ExportedProgramPassBase):
+
+  def __init__(self):
+    super().__init__()
+    self._patterns = [
+        _get_upsample_bilinear2d_pattern(),
+        _get_upsample_bilinear2d_align_corners_pattern(),
+        _get_interpolate_nearest2d_pattern(),
+    ]
+
+  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 ExportedProgramPassResult(exported_program, True)

ai_edge_torch/convert/fx_passes/canonicalize_pass.py

@@ -0,0 +1,37 @@
+# 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 torch
+from torch.export import ExportedProgram
+
+from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassBase
+from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassResult  # NOQA
+
+# 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,
+}
+
+
+class CanonicalizePass(ExportedProgramPassBase):
+
+  def call(self, exported_program: ExportedProgram):
+    exported_program = exported_program.run_decompositions(_dummy_decomp_table)
+    return ExportedProgramPassResult(exported_program, True)