ai-edge-torch-nightly 0.2.0.dev20240805__py3-none-any.whl → 0.2.0.dev20240807__py3-none-any.whl

ai_edge_torch/__init__.py

@@ -13,11 +13,11 @@
 # limitations under the License.
 # ==============================================================================
 
-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
-from .version import __version__
+from ai_edge_torch._convert.converter import convert
+from ai_edge_torch._convert.converter import signature
+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:

ai_edge_torch/{convert → _convert}/conversion.py

@@ -13,48 +13,44 @@
 # limitations under the License.
 # ==============================================================================
 
-import gc
 import logging
 import os
-from typing import Optional
+from typing import Any, Optional
 
+from ai_edge_torch import lowertools
 from ai_edge_torch import model
-from ai_edge_torch.convert import conversion_utils as cutils
-from ai_edge_torch.convert.fx_passes import BuildAtenCompositePass
-from ai_edge_torch.convert.fx_passes import BuildInterpolateCompositePass  # NOQA
-from ai_edge_torch.convert.fx_passes import CanonicalizePass
-from ai_edge_torch.convert.fx_passes import InjectMlirDebuginfoPass
-from ai_edge_torch.convert.fx_passes import OptimizeLayoutTransposesPass
-from ai_edge_torch.convert.fx_passes import run_passes
-from ai_edge_torch.generative.fx_passes import run_generative_passes
+from ai_edge_torch._convert import fx_passes
+from ai_edge_torch._convert import signature
+from ai_edge_torch.generative import fx_passes as generative_fx_passes
 from ai_edge_torch.quantize import quant_config as qcfg
 import torch
-from torch.export import ExportedProgram
-from torch_xla import stablehlo
 
 os.environ["EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM"] = "1"
 
 
 def _run_convert_passes(
-    exported_program: ExportedProgram,
-) -> ExportedProgram:
-  exported_program = run_generative_passes(exported_program)
-  return run_passes(
+    exported_program: torch.export.ExportedProgram,
+) -> torch.export.ExportedProgram:
+  exported_program = generative_fx_passes.run_generative_passes(
+      exported_program
+  )
+  return fx_passes.run_passes(
       exported_program,
       [
-          BuildInterpolateCompositePass(),
-          CanonicalizePass(),
-          OptimizeLayoutTransposesPass(),
-          CanonicalizePass(),
-          BuildAtenCompositePass(),
-          CanonicalizePass(),
-          InjectMlirDebuginfoPass(),
-          CanonicalizePass(),
+          fx_passes.BuildInterpolateCompositePass(),
+          fx_passes.CanonicalizePass(),
+          fx_passes.OptimizeLayoutTransposesPass(),
+          fx_passes.CanonicalizePass(),
+          fx_passes.BuildAtenCompositePass(),
+          fx_passes.CanonicalizePass(),
+          fx_passes.InjectMlirDebuginfoPass(),
+          fx_passes.CanonicalizePass(),
       ],
   )
 
 
-def _warn_training_modules(signatures: list[cutils.Signature]):
+def _warn_training_modules(signatures: list[signature.Signature]):
+  """Warns the user if the module is in training mode (.eval not called)."""
   for sig in signatures:
     if not sig.module.training:
       continue
@@ -64,30 +60,39 @@ def _warn_training_modules(signatures: list[cutils.Signature]):
         " module in evaluation mode with `module.eval()` for better on-device"
         " performance and compatibility."
     )
-    if len(signatures) == 1 and sig.name == cutils.DEFAULT_SIGNATURE_NAME:
+    if len(signatures) == 1 and sig.name == model.DEFAULT_SIGNATURE_NAME:
       # User does not specify any signature names explicitly.
       message = message.format(sig_name="")
     else:
       message = message.format(sig_name=f'"{sig.name}" ')
 
-    logging.warn(message)
+    logging.warning(message)
 
 
 def convert_signatures(
-    signatures: list[cutils.Signature],
+    signatures: list[signature.Signature],
     *,
     quant_config: Optional[qcfg.QuantConfig] = None,
-    _tfl_converter_flags: dict = {},
+    _tfl_converter_flags: Optional[dict[str, Any]],
 ) -> model.TfLiteModel:
-  """Converts a list of `Signature`s and embeds them into one `model.TfLiteModel`.
+  """Converts a list of `signature.Signature`s and embeds them into one `model.TfLiteModel`.
+
   Args:
-      signatures: The list of 'Signature' objects containing PyTorch modules to be converted.
+      signatures: The list of 'signature.Signature' objects containing PyTorch
+        modules to be converted.
       quant_config: User-defined quantization method and scheme of the model.
-      _tfl_converter_flags: A nested dictionary allowing setting flags for the underlying tflite converter.
+      _tfl_converter_flags: A nested dictionary allowing setting flags for the
+        underlying tflite converter.
+
+  Returns:
+    The converted `model.TfLiteModel` object.
   """
+  if _tfl_converter_flags is None:
+    _tfl_converter_flags = {}
+
   _warn_training_modules(signatures)
 
-  exported_programs: torch.export.ExportedProgram = [
+  exported_programs: torch.export.torch.export.ExportedProgram = [
       torch.export.export(
           sig.module, sig.flat_args, dynamic_shapes=sig.dynamic_shapes
       )
@@ -96,23 +101,8 @@ def convert_signatures(
 
   # Apply default fx passes
   exported_programs = list(map(_run_convert_passes, exported_programs))
-  shlo_bundles: list[stablehlo.StableHLOModelBundle] = [
-      cutils.exported_program_to_stablehlo_bundle(exported, sig.flat_args)
-      for exported, sig in zip(exported_programs, signatures)
-  ]
-
-  merged_shlo_graph_module: stablehlo.StableHLOGraphModule = (
-      cutils.merge_stablehlo_bundles(
-          shlo_bundles, signatures, exported_programs
-      )
-  )
-  del exported_programs
-  del shlo_bundles
-
-  gc.collect()
-
-  tflite_model = cutils.convert_stablehlo_to_tflite(
-      merged_shlo_graph_module,
+  tflite_model = lowertools.exported_programs_to_tflite(
+      exported_programs,
       signatures,
       quant_config=quant_config,
       _tfl_converter_flags=_tfl_converter_flags,

ai_edge_torch/_convert/conversion_utils.py

@@ -0,0 +1,64 @@
+# 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 Any
+
+from ai_edge_torch.quantize import quant_config as qcfg
+import tensorflow as tf
+
+
+def apply_tfl_converter_flags(
+    converter: tf.lite.TFLiteConverter, tfl_converter_flags: dict[str, Any]
+):
+  """Applies TFLite converter flags to the converter.
+
+  Args:
+    converter: TFLite converter.
+    tfl_converter_flags: TFLite converter flags.
+  """
+
+  def _set_converter_flag(path: list[Any]):
+    if len(path) < 2:
+      raise ValueError("Expecting at least two values in the path.")
+
+    target_obj = converter
+    for idx in range(len(path) - 2):
+      target_obj = getattr(target_obj, path[idx])
+
+    setattr(target_obj, path[-2], path[-1])
+
+  def _iterate_dict_tree(flags_dict: dict[str, Any], path: list[Any]):
+    for key, value in flags_dict.items():
+      path.append(key)
+      if isinstance(value, dict):
+        _iterate_dict_tree(value, path)
+      else:
+        path.append(value)
+        _set_converter_flag(path)
+        path.pop()
+      path.pop()
+
+  _iterate_dict_tree(tfl_converter_flags, [])
+
+
+def set_tfl_converter_quant_flags(
+    converter: tf.lite.TFLiteConverter, quant_config: qcfg.QuantConfig
+):
+  if quant_config is not None:
+    quantizer_mode = quant_config._quantizer_mode
+    if quantizer_mode == qcfg.QuantConfig._QuantizerMode.PT2E_DYNAMIC:
+      converter._experimental_qdq_conversion_mode = "DYNAMIC"
+    elif quantizer_mode == qcfg.QuantConfig._QuantizerMode.PT2E_STATIC:
+      converter._experimental_qdq_conversion_mode = "STATIC"

ai_edge_torch/{convert → _convert}/converter.py

@@ -15,19 +15,23 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 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._convert import conversion
+from ai_edge_torch._convert import signature as signature_module
 from ai_edge_torch.quantize import quant_config as qcfg
 import torch
 
 
 class Converter:
+  """A converter for converting PyTorch models to edge models.
+
+  This class allows adding multiple signatures to the converted edge model.
+  """
 
   def __init__(self):
-    self._signatures: list[cutils.Signature] = []
+    self._signatures: list[signature_module.Signature] = []
 
   def signature(
       self,
@@ -36,9 +40,9 @@ class Converter:
       sample_args=None,
       sample_kwargs=None,
       *,
-      dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+      dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
   ) -> Converter:
-    """Alias to `add_signature`"""
+    """Functions as an alias to `add_signature`."""
     return self.add_signature(
         name, module, sample_args, sample_kwargs, dynamic_shapes=dynamic_shapes
     )
@@ -50,17 +54,24 @@ class Converter:
       sample_args=None,
       sample_kwargs=None,
       *,
-      dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = 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.
+      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.
+
+    Returns:
+      The converter object itself.
 
     Raises:
       ValueError: If a signature with the provided name already exists.
@@ -75,7 +86,7 @@ class Converter:
       raise ValueError("sample_args or sample_kwargs must be provided.")
 
     self._signatures.append(
-        cutils.Signature(
+        signature_module.Signature(
             name,
             module,
             sample_args,
@@ -92,8 +103,8 @@ class Converter:
       sample_kwargs=None,
       *,
       quant_config: Optional[qcfg.QuantConfig] = None,
-      dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
-      _ai_edge_converter_flags: dict = {},
+      dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
+      _ai_edge_converter_flags: Optional[dict[str, Any]] = None,
   ) -> model.TfLiteModel:
     """Finalizes the conversion and produces an edge model.
 
@@ -101,31 +112,44 @@ class Converter:
 
       edge_model = Converter().signature(name, module, args).convert()
 
-    Or it could be used to set the default signature for the converted edge model:
+    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.
+      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.
+      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.
+
+    Returns:
+      The converted edge model.
 
     Raises:
-      ValueError: If the arguments are not provided as expected. See the example in this functions's comment.
+      ValueError: If the arguments are not provided as expected. See the example
+      in this functions's comment.
     """
+    if _ai_edge_converter_flags is None:
+      _ai_edge_converter_flags = {}
+
     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,
+            model.DEFAULT_SIGNATURE_NAME,
             module,
             sample_args,
             sample_kwargs,
@@ -136,7 +160,6 @@ class Converter:
             "sample_args or sample_kwargs must be provided if a module is"
             " specified."
         )
-
     return conversion.convert_signatures(
         self._signatures,
         quant_config=quant_config,
@@ -149,22 +172,28 @@ def signature(
     module: torch.nn.Module,
     sample_args=None,
     sample_kwargs=None,
-    dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = 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.
+    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.
+
+  Returns:
+    A Converter object with the provided signature.
 
   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
@@ -177,27 +206,38 @@ def convert(
     sample_kwargs=None,
     *,
     quant_config: Optional[qcfg.QuantConfig] = None,
-    dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
-    _ai_edge_converter_flags: dict = {},
+    dynamic_shapes: Optional[Union[dict[str, Any], Tuple[Any, ...]]] = None,
+    _ai_edge_converter_flags: Optional[dict[str, Any]] = None,
 ) -> model.TfLiteModel:
-  """Allows converting a PyTorch model to an edge model with one default signature in one step.
+  """Converts a PyTorch model to an edge model with a default signature.
 
   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.
+    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.
+    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.
+
+  Returns:
+    The converted edge model.
 
   Example:
     edge_model = ai_edge_torch.convert(module, args)
-
   """
 
+  if _ai_edge_converter_flags is None:
+    _ai_edge_converter_flags = {}
+
   return Converter().convert(
       module,
       sample_args,

ai_edge_torch/{convert → _convert}/fx_passes/__init__.py

@@ -15,15 +15,15 @@
 
 from typing import Sequence, Union
 
-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
+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
 from torch.export import ExportedProgram
 from torch.fx.passes.infra.pass_manager import pass_result_wrapper
 import torch.utils._pytree as pytree

ai_edge_torch/{convert → _convert}/fx_passes/build_aten_composite_pass.py

@@ -13,27 +13,23 @@
 # limitations under the License.
 # ==============================================================================
 
-import copy
-import functools
 from functools import reduce
 from typing import Any, Callable
-
-from ai_edge_torch.hlfb import StableHLOCompositeBuilder
+from ai_edge_torch import lowertools
 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
+from torch.fx.passes.infra import pass_base
 import torch.utils._pytree as pytree
 
-_composite_builders: dict[Callable, Callable[[GraphModule, Node], None]] = {}
+_composite_builders: dict[
+    Callable, Callable[[torch.fx.GraphModule, torch.fx.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
+      for overload in op.overloads():
+        _composite_builders[getattr(op, overload)] = func
     else:
       _composite_builders[op] = func
     return func
@@ -44,6 +40,19 @@ def _register_composite_builder(op):
 def _tree_map_to_composite_attr_values(
     values, *, stringify_incompatible_values=True
 ):
+  """Convert a tree of values to a tree of composite attribute values.
+
+  This is used for pre-processing op attributes before passing them to
+  the composite op as attributes.
+
+  Args:
+    values: A tree of values.
+    stringify_incompatible_values: If True, stringify values that are not
+      compatible with composite attributes.
+
+  Returns:
+    A tree of composite attribute values.
+  """
 
   def convert(value):
     nonlocal stringify_incompatible_values
@@ -60,6 +69,11 @@ def _tree_map_to_composite_attr_values(
 
 
 class TorchOpArgumentsMapper:
+  """A helper class to map op arguments to kwargs.
+
+  This is mainly used to extract the default values for op arguments and present
+  all arguments as kwargs.
+  """
 
   def __init__(self, op):
     if isinstance(op, torch._ops.OpOverloadPacket):
@@ -72,13 +86,21 @@ class TorchOpArgumentsMapper:
     ]
 
   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.
+    """Extracts all arguments of the op as kwargs.
+
+    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.
+
+    Args:
+      args: The op's arguments.
+      kwargs: The op's kwargs.
+
+    Returns:
+      A kwargs dict with all args and kwargs.
     """
     full_kwargs = {**(kwargs or {})}
 
-    for arg, (name, default_value) in zip(args, self.arg_specs):
+    for arg, (name, _) in zip(args, self.arg_specs):
       full_kwargs[name] = arg
 
     for name, default_value in self.arg_specs[len(args) :]:
@@ -89,12 +111,13 @@ class TorchOpArgumentsMapper:
 
 
 @_register_composite_builder(torch.ops.aten.hardswish.default)
-def _aten_hardswish(gm: GraphModule, node: Node):
+def _aten_hardswish(_: torch.fx.GraphModule, node: torch.fx.Node):
+  """Build a composite for aten.hardswish.default."""
   op = node.target
 
   def hardswish(self: torch.Tensor):
     nonlocal op
-    builder = StableHLOCompositeBuilder("aten.hardswish.default")
+    builder = lowertools.StableHLOCompositeBuilder("aten.hardswish.default")
     self = builder.mark_inputs(self)
     output = op(self)
     output = builder.mark_outputs(output)
@@ -104,7 +127,8 @@ def _aten_hardswish(gm: GraphModule, node: Node):
 
 
 @_register_composite_builder(torch.ops.aten.gelu.default)
-def _aten_gelu(gm: GraphModule, node: Node):
+def _aten_gelu(_: torch.fx.GraphModule, node: torch.fx.Node):
+  """Build a composite for aten.gelu.default."""
   op = node.target
   args_mapper = TorchOpArgumentsMapper(op)
 
@@ -120,7 +144,7 @@ def _aten_gelu(gm: GraphModule, node: Node):
     ):
       return op(*args, **kwargs)
 
-    builder = StableHLOCompositeBuilder(
+    builder = lowertools.StableHLOCompositeBuilder(
         "aten.gelu.default",
         attr=_tree_map_to_composite_attr_values({
             "approximate": full_kwargs["approximate"],
@@ -135,7 +159,8 @@ def _aten_gelu(gm: GraphModule, node: Node):
 
 
 @_register_composite_builder(torch.ops.aten.avg_pool2d.default)
-def _aten_avg_pool2d(gm: GraphModule, node: Node):
+def _aten_avg_pool2d(_: torch.fx.GraphModule, node: torch.fx.Node):
+  """Build a composite for aten.avg_pool2d.default."""
   op = node.target
   args_mapper = TorchOpArgumentsMapper(op)
 
@@ -199,7 +224,7 @@ def _aten_avg_pool2d(gm: GraphModule, node: Node):
     ):
       return op(*args, **kwargs)
 
-    builder = StableHLOCompositeBuilder(
+    builder = lowertools.StableHLOCompositeBuilder(
         "aten.avg_pool2d.default",
         attr=_tree_map_to_composite_attr_values({
             "kernel_size": full_kwargs["kernel_size"],
@@ -220,7 +245,7 @@ def _aten_avg_pool2d(gm: GraphModule, node: Node):
 
 
 @_register_composite_builder(torch.ops.aten.embedding.default)
-def _aten_embedding(gm: GraphModule, node: Node):
+def _aten_embedding(gm: torch.fx.GraphModule, node: torch.fx.Node):
   op = node.target
   args_mapper = TorchOpArgumentsMapper(op)
 
@@ -237,7 +262,7 @@ def _aten_embedding(gm: GraphModule, node: Node):
     # Explicitly reshape to 1D. This places the ReshapeOp outside of the HLFB.
     idx = torch.reshape(idx, (idx.numel(),))
 
-    builder = StableHLOCompositeBuilder("odml.embedding_lookup")
+    builder = lowertools.StableHLOCompositeBuilder("odml.embedding_lookup")
     full_kwargs["indices"], full_kwargs["weight"] = builder.mark_inputs(
         idx,
         full_kwargs["weight"],
@@ -252,13 +277,13 @@ def _aten_embedding(gm: GraphModule, node: Node):
   node.target = embedding
 
 
-class BuildAtenCompositePass(PassBase):
+class BuildAtenCompositePass(pass_base.PassBase):
 
-  def call(self, graph_module: GraphModule):
+  def call(self, graph_module: torch.fx.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)
+    return pass_base.PassResult(graph_module, True)