ai-edge-torch-nightly 0.2.0.dev20240801__py3-none-any.whl → 0.2.0.dev20240803__py3-none-any.whl

ai_edge_torch/__init__.py

@@ -17,6 +17,7 @@ from .convert.converter import convert
 from .convert.converter import signature
 from .convert.to_channel_last_io import to_channel_last_io
 from .model import Model
+from .version import __version__
 
 
 def load(path: str) -> Model:

ai_edge_torch/convert/conversion.py

@@ -18,10 +18,6 @@ import logging
 import os
 from typing import Optional
 
-import torch
-from torch.export import ExportedProgram
-from torch_xla import stablehlo
-
 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
@@ -32,6 +28,9 @@ 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.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"
 
@@ -61,8 +60,9 @@ def _warn_training_modules(signatures: list[cutils.Signature]):
       continue
 
     message = (
-        "Your model {sig_name}is converted in training mode. "
-        "Please set the module in evaluation mode with `module.eval()` for better on-device performance and compatibility."
+        "Your model {sig_name}is converted in training mode. Please set the"
+        " 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:
       # User does not specify any signature names explicitly.
@@ -88,7 +88,9 @@ def convert_signatures(
   _warn_training_modules(signatures)
 
   exported_programs: torch.export.ExportedProgram = [
-      torch.export.export(sig.module, sig.flat_args, dynamic_shapes=sig.dynamic_shapes)
+      torch.export.export(
+          sig.module, sig.flat_args, dynamic_shapes=sig.dynamic_shapes
+      )
       for sig in signatures
   ]
 
@@ -100,7 +102,9 @@ def convert_signatures(
   ]
 
   merged_shlo_graph_module: stablehlo.StableHLOGraphModule = (
-      cutils.merge_stablehlo_bundles(shlo_bundles, signatures, exported_programs)
+      cutils.merge_stablehlo_bundles(
+          shlo_bundles, signatures, exported_programs
+      )
   )
   del exported_programs
   del shlo_bundles

ai_edge_torch/convert/conversion_utils.py

@@ -22,15 +22,15 @@ import logging
 import tempfile
 from typing import Any, Dict, List, Optional, Tuple, Union
 
+from ai_edge_torch.generative.quantize.ai_edge_quantizer_glue import translate_recipe  # NOQA
+from ai_edge_torch.quantize import quant_config as qcfg
 import torch
 import torch.utils._pytree as pytree
 from torch_xla import stablehlo
 
-from ai_edge_torch.generative.quantize.ai_edge_quantizer_glue import translate_recipe  # NOQA
-from ai_edge_torch.quantize import quant_config as qcfg
-
 try:
   import tensorflow as tf
+
   from tensorflow.compiler.tf2xla.python import xla as tfxla
 
   from tensorflow.lite.python import conversion_metadata_schema_py_generated as conversion_metadata_fb  # isort:skip
@@ -90,18 +90,20 @@ class Signature:
     if context is None:
       for i, spec in enumerate(specs):
         if spec.children_specs:
-          flat_names.extend(
-              [
-                  f"{i}_{name}"
-                  for name in self._flat_kwarg_names(spec.children_specs, spec.context)
-              ]
-          )
+          flat_names.extend([
+              f"{i}_{name}"
+              for name in self._flat_kwarg_names(
+                  spec.children_specs, spec.context
+              )
+          ])
         else:
           flat_names.append(f"{i}")
     else:
       flat_ctx = self._flatten_list(context)
       for prefix, spec in zip(flat_ctx, specs):
-        leaf_flat_names = self._flat_kwarg_names(spec.children_specs, spec.context)
+        leaf_flat_names = self._flat_kwarg_names(
+            spec.children_specs, spec.context
+        )
         if leaf_flat_names:
           flat_names.extend([f"{prefix}_{name}" for name in leaf_flat_names])
         else:
@@ -125,7 +127,8 @@ class Signature:
 
 
 def exported_program_to_stablehlo_bundle(
-    exported_program: torch.export.ExportedProgram, sample_args: tuple[torch.Tensor]
+    exported_program: torch.export.ExportedProgram,
+    sample_args: tuple[torch.Tensor],
 ) -> stablehlo.StableHLOModelBundle:
   # Setting export_weights to False here so that pytorch/xla avoids copying the weights
   # to a numpy array which would lead to memory bloat. This means that the state_dict
@@ -146,7 +149,9 @@ def _torch_to_tf_tensor(torch_tensor: torch.Tensor):
     dlpack_capsule = torch.utils.dlpack.to_dlpack(torch_tensor)
     tf_tensor = tf.experimental.dlpack.from_dlpack(dlpack_capsule)
   except Exception:
-    logging.info("Can not use dlpack to convert torch tensors. Falling back to numpy.")
+    logging.info(
+        "Can not use dlpack to convert torch tensors. Falling back to numpy."
+    )
     nparray = torch_tensor.cpu().detach().numpy()
     tf_tensor = tf.convert_to_tensor(nparray)
 
@@ -154,7 +159,8 @@ def _torch_to_tf_tensor(torch_tensor: torch.Tensor):
 
 
 def _get_states(
-    exported_programs: list[torch.export.ExportedProgram], signatures: list[Signature]
+    exported_programs: list[torch.export.ExportedProgram],
+    signatures: list[Signature],
 ):
   for exported_program, signature in zip(exported_programs, signatures):
     args, _ = exported_program.example_inputs
@@ -166,7 +172,8 @@ def _get_states(
       # Only interested in Tensors that are part of the state (and not user input).
       if (
           not isinstance(tensor, torch.Tensor)
-          or input_spec.kind == torch.export.graph_signature.InputKind.USER_INPUT
+          or input_spec.kind
+          == torch.export.graph_signature.InputKind.USER_INPUT
       ):
         continue
       yield signature, tensor, input_spec
@@ -192,9 +199,13 @@ def _gather_state_dict(
     deduped_tensor_map[unique_id] = _torch_to_tf_tensor(tensor)
 
   state_dict = {}
-  for signature, tensor, input_spec in _get_states(exported_programs, signatures):
+  for signature, tensor, input_spec in _get_states(
+      exported_programs, signatures
+  ):
     unique_id = _tensor_unique_id(tensor)
-    state_dict[signature.name + "_" + input_spec.target] = deduped_tensor_map[unique_id]
+    state_dict[signature.name + "_" + input_spec.target] = deduped_tensor_map[
+        unique_id
+    ]
 
   return state_dict
 
@@ -236,7 +247,9 @@ def _wrap_as_tf_func(
 ):
   def inner(*args):
     type_info = [sig.dtype for sig in func.meta.output_signature]
-    shape_info = [_get_shape_with_dynamic(sig) for sig in func.meta.output_signature]
+    shape_info = [
+        _get_shape_with_dynamic(sig) for sig in func.meta.output_signature
+    ]
     call_args = stablehlo._extract_call_parameters(args, func.meta, bundle)
     return tfxla.call_module(
         tuple(call_args),
@@ -369,7 +382,9 @@ def convert_stablehlo_to_tflite(
         )
     )
 
-  tf_module._variables = list(bundle.state_dict.values()) + bundle.additional_constants
+  tf_module._variables = (
+      list(bundle.state_dict.values()) + bundle.additional_constants
+  )
   del bundle
   gc.collect()
 
@@ -385,7 +400,8 @@ def convert_stablehlo_to_tflite(
         tf_module,
         temp_dir_path,
         signatures={
-            sig.name: tf_concrete_funcs[idx] for idx, sig in enumerate(signatures)
+            sig.name: tf_concrete_funcs[idx]
+            for idx, sig in enumerate(signatures)
         },
     )
     # Clean up intermediate memory early.
@@ -416,6 +432,8 @@ def convert_stablehlo_to_tflite(
         and quant_config._quantizer_mode
         == quant_config._QuantizerMode.AI_EDGE_QUANTIZER
     ):
-      tflite_model = translate_recipe.quantize_model(tflite_model, translated_recipe)
+      tflite_model = translate_recipe.quantize_model(
+          tflite_model, translated_recipe
+      )
 
   return tflite_model

ai_edge_torch/convert/converter.py

@@ -17,12 +17,11 @@ 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
+import torch
 
 
 class Converter:
@@ -68,14 +67,20 @@ class Converter:
     """
 
     if name in [sig.name for sig in self._signatures]:
-      raise ValueError(f"A signature with the provided name ({name}) is already added.")
+      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
+            name,
+            module,
+            sample_args,
+            sample_kwargs,
+            dynamic_shapes=dynamic_shapes,
         )
     )
     return self
@@ -128,7 +133,8 @@ class Converter:
         )
       else:  # module is provided but not args
         raise ValueError(
-            "sample_args or sample_kwargs must be provided if a module is specified."
+            "sample_args or sample_kwargs must be provided if a module is"
+            " specified."
         )
 
     return conversion.convert_signatures(

ai_edge_torch/convert/fx_passes/__init__.py

@@ -15,10 +15,6 @@
 
 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
@@ -28,6 +24,9 @@ from ai_edge_torch.convert.fx_passes.build_interpolate_composite_pass import Bui
 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
 
 
 # TODO(cnchan): make a PassManager class.

ai_edge_torch/convert/fx_passes/_pass_base.py

@@ -32,14 +32,18 @@ class ExportedProgramPassResult(
 
 class ExportedProgramPassBase(abc.ABC):
 
-  def __call__(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
+  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:
+  def call(
+      self, exported_program: ExportedProgram
+  ) -> ExportedProgramPassResult:
     pass
 
   def requires(self, exported_program: ExportedProgram) -> None:

ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py

@@ -15,8 +15,10 @@
 
 import copy
 import functools
+from functools import reduce
 from typing import Any, Callable
 
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
 import torch
 from torch.fx import GraphModule
 from torch.fx import Node
@@ -24,8 +26,6 @@ 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]] = {}
 
 
@@ -41,7 +41,9 @@ def _register_composite_builder(op):
   return inner
 
 
-def _tree_map_to_composite_attr_values(values, *, stringify_incompatible_values=True):
+def _tree_map_to_composite_attr_values(
+    values, *, stringify_incompatible_values=True
+):
 
   def convert(value):
     nonlocal stringify_incompatible_values
@@ -65,7 +67,9 @@ class TorchOpArgumentsMapper:
 
     assert hasattr(op, "_schema")
     self.op = op
-    self.arg_specs = [(spec.name, spec.default_value) for spec in op._schema.arguments]
+    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
@@ -110,16 +114,17 @@ def _aten_gelu(gm: GraphModule, node: Node):
     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":
+    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"],
-            }
-        ),
+        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"])
@@ -150,7 +155,10 @@ def _aten_avg_pool2d(gm: GraphModule, node: Node):
       ):
         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
+            0,
+            (dim_output_size - 1) * dim_stride
+            + dim_kernel_size
+            - dim_input_size,
         )
         if padding_needed % 2 != 0:
           return False
@@ -193,16 +201,14 @@ def _aten_avg_pool2d(gm: GraphModule, node: Node):
 
     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"],
-            }
-        ),
+        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"])
@@ -223,25 +229,25 @@ def _aten_embedding(gm: GraphModule, node: Node):
     full_kwargs = args_mapper.get_full_kwargs(args, kwargs)
     _, embedding_dim = full_kwargs["weight"].size()
     idx = full_kwargs["indices"]
-    # TODO(b/356458830): Handle relative positional encoding
-    if len(idx.size()) == 2:
-      idx = idx.type(torch.int)
-      B, T = idx.size()
-
-      idx = torch.reshape(idx, (B * T,))
-
-      builder = StableHLOCompositeBuilder("odml.embedding_lookup")
-      full_kwargs["indices"], full_kwargs["weight"] = builder.mark_inputs(
-          idx,
-          full_kwargs["weight"],
-      )
-      output = op(**full_kwargs)
-      output = builder.mark_outputs(output)
-
-      output = torch.reshape(output, (B, T, embedding_dim))
-      return output
-    else:
-      return op(**full_kwargs)
+
+    # Explicitly cast to INT32. This places the CastOp outside of the HLFB.
+    idx = idx.type(torch.int)
+    original_idx_shape = idx.size()
+
+    # Explicitly reshape to 1D. This places the ReshapeOp outside of the HLFB.
+    idx = torch.reshape(idx, (idx.numel(),))
+
+    builder = StableHLOCompositeBuilder("odml.embedding_lookup")
+    full_kwargs["indices"], full_kwargs["weight"] = builder.mark_inputs(
+        idx,
+        full_kwargs["weight"],
+    )
+    output = op(**full_kwargs)
+    output = builder.mark_outputs(output)
+
+    # Explicitly reshape back to the original shape. This places the ReshapeOp outside of the HLFB.
+    output = torch.reshape(output, (*(original_idx_shape), embedding_dim))
+    return output
 
   node.target = embedding
 

ai_edge_torch/convert/fx_passes/build_interpolate_composite_pass.py

@@ -15,23 +15,20 @@
 
 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
+import torch
 
 # 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,
-    ]
-)
+_INTERPOLATE_DECOMPOSITIONS = torch._decomp.get_decompositions([
+    torch.ops.aten.upsample_bilinear2d.vec,
+    torch.ops.aten.upsample_nearest2d.vec,
+])
 
 
 @functools.cache
@@ -84,7 +81,9 @@ def _get_upsample_bilinear2d_align_corners_pattern():
 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"),
+      lambda x: torch.nn.functional.interpolate(
+          x, scale_factor=2, mode="nearest"
+      ),
       export_args=(torch.rand(1, 3, 100, 100),),
       decomp_table=_INTERPOLATE_DECOMPOSITIONS,
   )
@@ -112,7 +111,9 @@ class BuildInterpolateCompositePass(ExportedProgramPassBase):
     ]
 
   def call(self, exported_program: torch.export.ExportedProgram):
-    exported_program = exported_program.run_decompositions(_INTERPOLATE_DECOMPOSITIONS)
+    exported_program = exported_program.run_decompositions(
+        _INTERPOLATE_DECOMPOSITIONS
+    )
 
     graph_module = exported_program.graph_module
     for pattern in self._patterns:

ai_edge_torch/convert/fx_passes/canonicalize_pass.py

@@ -13,11 +13,10 @@
 # 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
+import torch
+from torch.export import ExportedProgram
 
 # A dummy decomp table for running ExportedProgram.run_decompositions without
 # any op decompositions but just aot_export_module. Due to the check in

ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_check.py

@@ -15,13 +15,12 @@
 import dataclasses
 import operator
 
-import torch
-from torch.fx import Node
-
 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_rewrite  # NOQA
 from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import utils  # NOQA
 from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass.op_func_registry import OpFuncRegistry  # NOQA
+import torch
+from torch.fx import Node
 
 aten = torch.ops.aten
 
@@ -150,7 +149,9 @@ def _qdq_layout_sensitive_inputs_getter(node: Node):
 
 
 @layout_sensitive_inputs_getters.register(aten.convolution)
-@layout_sensitive_inputs_getters.register(aten._native_batch_norm_legit_no_training)
+@layout_sensitive_inputs_getters.register(
+    aten._native_batch_norm_legit_no_training
+)
 @layout_sensitive_inputs_getters.register(aten.native_group_norm)
 def _first_arg_getter(node):
   return [node.args[0]]
@@ -174,7 +175,11 @@ def _all_layout_sensitive_inputs_are_4d_checker(node: Node):
 @nhwcable_node_checkers.register(aten._native_batch_norm_legit_no_training)
 def _aten_norm_checker(node):
   val = node.meta.get("val")
-  if not isinstance(val, (list, tuple)) or not val or not hasattr(val[0], "shape"):
+  if (
+      not isinstance(val, (list, tuple))
+      or not val
+      or not hasattr(val[0], "shape")
+  ):
     return NHWCable(can_be=False, must_be=False)
   return NHWCable(can_be=len(val[0].shape) == 4, must_be=False)
 
@@ -182,9 +187,15 @@ def _aten_norm_checker(node):
 @nhwcable_node_checkers.register(aten.native_group_norm)
 def _aten_native_group_norm_checker(node):
   val = node.meta.get("val")
-  if not isinstance(val, (list, tuple)) or not val or not hasattr(val[0], "shape"):
+  if (
+      not isinstance(val, (list, tuple))
+      or not val
+      or not hasattr(val[0], "shape")
+  ):
     return NHWCable(can_be=False, must_be=False)
-  if len(node.args) >= 3 and (node.args[1] is not None or node.args[2] is not None):
+  if len(node.args) >= 3 and (
+      node.args[1] is not None or node.args[2] is not None
+  ):
     # Disable NHWC rewriter due to precision issue with weight and bias.
     # TODO(b/354780253): Re-enable NHWC rewriter with proper lowering.
     return NHWCable(can_be=False, must_be=False)

ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py

@@ -13,10 +13,9 @@
 # limitations under the License.
 # ==============================================================================
 
-import torch
-
 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
+import torch
 
 
 def partition(graph_module: torch.fx.GraphModule):
@@ -45,7 +44,9 @@ def partition(graph_module: torch.fx.GraphModule):
 
       layout_sensitive_inputs = layout_check.get_layout_sensitive_inputs(node)
 
-      should_be_nhwc = any(map(layout_mark.is_nhwc_node, layout_sensitive_inputs))
+      should_be_nhwc = any(
+          map(layout_mark.is_nhwc_node, layout_sensitive_inputs)
+      )
       for input_node in layout_sensitive_inputs:
         if not layout_mark.is_nhwc_node(input_node) and not layout_check.is_4d(
             input_node

ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py

@@ -17,13 +17,12 @@ import collections
 import dataclasses
 import itertools
 
+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
 import numpy as np
 import scipy
 import torch
 
-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
-
 
 def can_partition(graph_module: torch.fx.GraphModule):
   """Returns true if the input graph_module can be partitioned by min cut solver
@@ -83,7 +82,10 @@ class MinCutSolver:
   def graph(self):
     edges = np.array(self.edges)
     return scipy.sparse.csr_matrix(
-        (np.minimum(edges[:, 2], MinCutSolver.INF_COST), (edges[:, 0], edges[:, 1])),
+        (
+            np.minimum(edges[:, 2], MinCutSolver.INF_COST),
+            (edges[:, 0], edges[:, 1]),
+        ),
         shape=(self._nodes_cnt, self._nodes_cnt),
         dtype=np.int32,
     )

ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py

@@ -14,13 +14,12 @@
 # ==============================================================================
 import operator
 
-import torch
-from torch.fx import Node
-import torch.utils._pytree as pytree
-
 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 utils  # NOQA
 from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass.op_func_registry import OpFuncRegistry  # NOQA
+import torch
+from torch.fx import Node
+import torch.utils._pytree as pytree
 
 aten = torch.ops.aten
 
@@ -349,7 +348,12 @@ def _aten_native_group_norm(node):
   ):
     input_reshaped = torch.reshape(
         input,
-        [batch_size, flattened_inner_size, num_groups, num_channels // num_groups],
+        [
+            batch_size,
+            flattened_inner_size,
+            num_groups,
+            num_channels // num_groups,
+        ],
     )
     reduction_dims = [1, 3]
 

ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/op_func_registry.py

@@ -12,9 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-import torch
-
 from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import utils  # NOQA
+import torch
 
 
 class OpFuncRegistry(dict):