ai-edge-torch-nightly 0.2.0.dev20240730__py3-none-any.whl → 0.2.0.dev20240802__py3-none-any.whl

ai_edge_torch/generative/utilities/t5_loader.py

@@ -18,11 +18,10 @@ import glob
 import os
 from typing import Callable, Dict
 
+from ai_edge_torch.generative.layers import model_config
 from safetensors import safe_open
 import torch
 
-from ai_edge_torch.generative.layers import model_config
-
 
 def load_safetensors(full_path: str):
   """Loads safetensors into a single state dictionary.
@@ -71,7 +70,11 @@ def load_pytorch_statedict(full_path: str):
   Raises:
     ValueError: If no tensors are loaded from the provided directory or file.
   """
-  pattern = os.path.join(full_path, "*.bin") if os.path.isdir(full_path) else full_path
+  pattern = (
+      os.path.join(full_path, "*.bin")
+      if os.path.isdir(full_path)
+      else full_path
+  )
   files = []
   for file in glob.glob(pattern):
     files.append(file)
@@ -131,7 +134,10 @@ class ModelLoader:
     self._loader = self._get_loader()
 
   def load(
-      self, model: torch.nn.Module, strict: bool = True, fuse_attention: bool = True
+      self,
+      model: torch.nn.Module,
+      strict: bool = True,
+      fuse_attention: bool = True,
   ):
     """Load the model from the checkpoint
 
@@ -166,11 +172,14 @@ class ModelLoader:
 
     if strict and state:
       raise ValueError(
-          f"Failed to map all tensor. Remaining tensor are: {list(state.keys())}"
+          "Failed to map all tensor. Remaining tensor are:"
+          f" {list(state.keys())}"
       )
     model.load_state_dict(converted_state, strict=strict)
 
-  def _do_load(self, model, state, names, additional_prefix="", fuse_attention=True):
+  def _do_load(
+      self, model, state, names, additional_prefix="", fuse_attention=True
+  ):
     """Load the model from the checkpoint
 
     Args:
@@ -183,7 +192,9 @@ class ModelLoader:
     """
     converted_state = dict()
     if names.embedding is not None:
-      converted_state["tok_embedding.weight"] = state.pop(f"{names.embedding}.weight")
+      converted_state["tok_embedding.weight"] = state.pop(
+          f"{names.embedding}.weight"
+      )
     if names.lm_head is not None:
       converted_state["lm_head.weight"] = state.pop(f"{names.lm_head}.weight")
       if model.config.lm_head_use_bias:
@@ -195,7 +206,9 @@ class ModelLoader:
           f"{final_norm_name}.weight"
       )
       if f"{final_norm_name}.bias" in state:
-        converted_state["final_norm.bias"] = state.pop(f"{final_norm_name}.bias")
+        converted_state["final_norm.bias"] = state.pop(
+            f"{final_norm_name}.bias"
+        )
 
     if names.relative_attn_bias:
       rel_attn_name = names.relative_attn_bias
@@ -205,7 +218,9 @@ class ModelLoader:
       )
 
     for i in range(model.config.num_layers):
-      self._map_norm(i, model.config, state, converted_state, names, additional_prefix)
+      self._map_norm(
+          i, model.config, state, converted_state, names, additional_prefix
+      )
       self._map_feedforward(
           i, model.config, state, converted_state, names, additional_prefix
       )
@@ -268,13 +283,19 @@ class ModelLoader:
     if config.ff_config.type == model_config.FeedForwardType.SEQUENTIAL:
       ff_up_proj_name = names.ff_up_proj.format(idx)
       ff_down_proj_name = names.ff_down_proj.format(idx)
-      converted_state[f"{prefix}.ff.w1.weight"] = state.pop(f"{ff_up_proj_name}.weight")
+      converted_state[f"{prefix}.ff.w1.weight"] = state.pop(
+          f"{ff_up_proj_name}.weight"
+      )
       converted_state[f"{prefix}.ff.w2.weight"] = state.pop(
           f"{ff_down_proj_name}.weight"
       )
       if config.ff_config.use_bias:
-        converted_state[f"{prefix}.ff.w1.bias"] = state.pop(f"{ff_up_proj_name}.bias")
-        converted_state[f"{prefix}.ff.w2.bias"] = state.pop(f"{ff_down_proj_name}.bias")
+        converted_state[f"{prefix}.ff.w1.bias"] = state.pop(
+            f"{ff_up_proj_name}.bias"
+        )
+        converted_state[f"{prefix}.ff.w2.bias"] = state.pop(
+            f"{ff_down_proj_name}.bias"
+        )
     else:
       if names.ff_gate_proj is not None:
         ff_up_proj_name = names.ff_up_proj.format(idx)
@@ -290,7 +311,9 @@ class ModelLoader:
             f"{ff_gate_proj_name}.weight"
         )
         if config.ff_config.use_bias:
-          converted_state[f"{prefix}.ff.w3.bias"] = state.pop(f"{ff_up_proj_name}.bias")
+          converted_state[f"{prefix}.ff.w3.bias"] = state.pop(
+              f"{ff_up_proj_name}.bias"
+          )
           converted_state[f"{prefix}.ff.w2.bias"] = state.pop(
               f"{ff_down_proj_name}.bias"
           )
@@ -355,12 +378,12 @@ class ModelLoader:
         )
 
     o_name = names.attn_output_proj.format(idx)
-    converted_state[f"{prefix}.atten_func.output_projection.weight"] = state.pop(
-        f"{o_name}.weight"
+    converted_state[f"{prefix}.atten_func.output_projection.weight"] = (
+        state.pop(f"{o_name}.weight")
     )
     if config.attn_config.output_proj_use_bias:
-      converted_state[f"{prefix}.atten_func.output_projection.bias"] = state.pop(
-          f"{o_name}.bias"
+      converted_state[f"{prefix}.atten_func.output_projection.bias"] = (
+          state.pop(f"{o_name}.bias")
       )
 
   def _map_cross_attention(
@@ -385,47 +408,51 @@ class ModelLoader:
     v_name = names.cross_attn_value_proj.format(idx)
 
     if fuse_attention:
-      converted_state[f"{prefix}.cross_atten_func.attn.weight"] = self._fuse_qkv(
-          config,
-          state.pop(f"{q_name}.weight"),
-          state.pop(f"{k_name}.weight"),
-          state.pop(f"{v_name}.weight"),
+      converted_state[f"{prefix}.cross_atten_func.attn.weight"] = (
+          self._fuse_qkv(
+              config,
+              state.pop(f"{q_name}.weight"),
+              state.pop(f"{k_name}.weight"),
+              state.pop(f"{v_name}.weight"),
+          )
       )
       if config.attn_config.qkv_use_bias:
-        converted_state[f"{prefix}.cross_atten_func.attn.bias"] = self._fuse_qkv(
-            config,
-            state.pop(f"{q_name}.bias"),
-            state.pop(f"{k_name}.bias"),
-            state.pop(f"{v_name}.bias"),
+        converted_state[f"{prefix}.cross_atten_func.attn.bias"] = (
+            self._fuse_qkv(
+                config,
+                state.pop(f"{q_name}.bias"),
+                state.pop(f"{k_name}.bias"),
+                state.pop(f"{v_name}.bias"),
+            )
         )
     else:
-      converted_state[f"{prefix}.cross_atten_func.q_projection.weight"] = state.pop(
-          f"{q_name}.weight"
+      converted_state[f"{prefix}.cross_atten_func.q_projection.weight"] = (
+          state.pop(f"{q_name}.weight")
       )
-      converted_state[f"{prefix}.cross_atten_func.k_projection.weight"] = state.pop(
-          f"{k_name}.weight"
+      converted_state[f"{prefix}.cross_atten_func.k_projection.weight"] = (
+          state.pop(f"{k_name}.weight")
       )
-      converted_state[f"{prefix}.cross_atten_func.v_projection.weight"] = state.pop(
-          f"{v_name}.weight"
+      converted_state[f"{prefix}.cross_atten_func.v_projection.weight"] = (
+          state.pop(f"{v_name}.weight")
       )
       if config.attn_config.qkv_use_bias:
-        converted_state[f"{prefix}.cross_atten_func.q_projection.bias"] = state.pop(
-            f"{q_name}.bias"
+        converted_state[f"{prefix}.cross_atten_func.q_projection.bias"] = (
+            state.pop(f"{q_name}.bias")
         )
-        converted_state[f"{prefix}.cross_atten_func.k_projection.bias"] = state.pop(
-            f"{k_name}.bias"
+        converted_state[f"{prefix}.cross_atten_func.k_projection.bias"] = (
+            state.pop(f"{k_name}.bias")
         )
-        converted_state[f"{prefix}.cross_atten_func.v_projection.bias"] = state.pop(
-            f"{v_name}.bias"
+        converted_state[f"{prefix}.cross_atten_func.v_projection.bias"] = (
+            state.pop(f"{v_name}.bias")
         )
 
     o_name = names.cross_attn_output_proj.format(idx)
-    converted_state[f"{prefix}.cross_atten_func.output_projection.weight"] = state.pop(
-        f"{o_name}.weight"
+    converted_state[f"{prefix}.cross_atten_func.output_projection.weight"] = (
+        state.pop(f"{o_name}.weight")
     )
     if config.attn_config.output_proj_use_bias:
-      converted_state[f"{prefix}.cross_atten_func.output_projection.bias"] = state.pop(
-          f"{o_name}.bias"
+      converted_state[f"{prefix}.cross_atten_func.output_projection.bias"] = (
+          state.pop(f"{o_name}.bias")
       )
 
   def _map_norm(
@@ -450,12 +477,12 @@ class ModelLoader:
 
     if names.pre_cross_attn_norm:
       pre_cross_attn_norm_name = names.pre_cross_attn_norm.format(idx)
-      converted_state[f"{prefix}.cross_atten_func.pre_atten_norm.weight"] = state.pop(
-          f"{pre_cross_attn_norm_name}.weight"
+      converted_state[f"{prefix}.cross_atten_func.pre_atten_norm.weight"] = (
+          state.pop(f"{pre_cross_attn_norm_name}.weight")
       )
       if f"{pre_cross_attn_norm_name}.bias" in state:
-        converted_state[f"{prefix}.cross_atten_func.pre_atten_norm.bias"] = state.pop(
-            f"{pre_cross_attn_norm_name}.bias"
+        converted_state[f"{prefix}.cross_atten_func.pre_atten_norm.bias"] = (
+            state.pop(f"{pre_cross_attn_norm_name}.bias")
         )
 
     if names.pre_ff_norm is not None:
@@ -475,7 +502,9 @@ class ModelLoader:
       k: torch.Tensor,
       v: torch.Tensor,
   ) -> torch.Tensor:
-    q_per_kv = config.attn_config.num_heads // config.attn_config.num_query_groups
+    q_per_kv = (
+        config.attn_config.num_heads // config.attn_config.num_query_groups
+    )
     qs = torch.split(q, config.head_dim * q_per_kv)
     ks = torch.split(k, config.head_dim)
     vs = torch.split(v, config.head_dim)

ai_edge_torch/hlfb/mark_pattern/__init__.py

@@ -16,11 +16,10 @@ import copy
 from typing import Any
 import uuid
 
-import torch
-from torch_xla.experimental import xla_marker
-
 from ai_edge_torch.hlfb.mark_pattern.pattern import Pattern
 from ai_edge_torch.hlfb.mark_pattern.pattern import ScalarAttrTracker  # NOQA
+import torch
+from torch_xla.experimental import xla_marker
 
 
 @torch._dynamo.assume_constant_result

ai_edge_torch/hlfb/mark_pattern/pattern.py

@@ -16,6 +16,7 @@ import copy
 import dataclasses
 from typing import Any, Callable, Optional, Union
 
+from ai_edge_torch.hlfb.mark_pattern import passes
 import torch
 from torch.export.graph_signature import TensorArgument
 from torch.fx import Graph
@@ -23,8 +24,6 @@ from torch.fx import GraphModule
 from torch.fx.passes.utils.matcher_utils import InternalMatch
 from torch.fx.passes.utils.matcher_utils import SubgraphMatcher
 
-from ai_edge_torch.hlfb.mark_pattern import passes
-
 
 def _are_equal(x: Any, y: Any) -> bool:
   if type(x) != type(y):
@@ -69,7 +68,9 @@ class ScalarAttrTracker:
   pattern_arg_pos: int
   transform: Callable = lambda x: x
   inverse_transform: Callable = lambda x: x
-  _source_targets: list[tuple[Any, Any]] = dataclasses.field(default_factory=list)
+  _source_targets: list[tuple[Any, Any]] = dataclasses.field(
+      default_factory=list
+  )
 
   def track(self, *sources):
     """Register magic values to track the (transformed) attr values in
@@ -78,7 +79,9 @@ class ScalarAttrTracker:
     for source in sources:
       target = self.transform(source)
       if not _are_equal(self.inverse_transform(target), source):
-        raise Exception(f"Invalid transform/inverse_transform for {self.attr_name}")
+        raise Exception(
+            f"Invalid transform/inverse_transform for {self.attr_name}"
+        )
       self._source_targets.append([source, target])
     return self
 
@@ -189,7 +192,9 @@ class Pattern:
 
     self.name = name
     self.attr_builder = attr_builder
-    self._scalar_attr_trackers = scalar_attr_trackers if scalar_attr_trackers else []
+    self._scalar_attr_trackers = (
+        scalar_attr_trackers if scalar_attr_trackers else []
+    )
 
     exported_program = torch.export.export(module, export_args)
     if decomp_table is not None:
@@ -201,7 +206,9 @@ class Pattern:
     self._scalar_attr_locations = []
     for tracker in self._scalar_attr_trackers:
       self._scalar_attr_locations.append(
-          _find_scalar_attr(module, export_args, tracker, decomp_table=decomp_table)
+          _find_scalar_attr(
+              module, export_args, tracker, decomp_table=decomp_table
+          )
       )
 
     # Sanitize graph_module for more precise pattern matching.
@@ -251,7 +258,9 @@ class Pattern:
         attrs = {}
 
       for loc in self._scalar_attr_locations:
-        attrs[loc.attr_name] = self._get_attr_value_from_pattern_match(match, loc)
+        attrs[loc.attr_name] = self._get_attr_value_from_pattern_match(
+            match, loc
+        )
 
       attrs = attrs if attrs else None
       match_with_attrs.append((match, attrs))

ai_edge_torch/hlfb/test/test_mark_pattern.py

@@ -15,11 +15,10 @@
 
 import unittest
 
+from ai_edge_torch.hlfb import mark_pattern
 import torch
 import torch_xla
 
-from ai_edge_torch.hlfb import mark_pattern
-
 
 def _export_stablehlo_mlir(model, args=None):
   if not isinstance(model, torch.export.ExportedProgram):
@@ -73,7 +72,9 @@ class TestMarkPattern(unittest.TestCase):
     mlir = _export_stablehlo_mlir(exported_program)
 
     self.assertEqual(mlir.count('stablehlo.composite "test.add"'), 2)
-    self.assertEqual(mlir.count('composite_attributes = {alias = "test.test_add"}'), 2)
+    self.assertEqual(
+        mlir.count('composite_attributes = {alias = "test.test_add"}'), 2
+    )
 
   def test_mark_pattern_with_scalar_attr_tracker(self):
     class TestModel(torch.nn.Module):

ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py

@@ -15,12 +15,11 @@
 import math
 import unittest
 
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
 import torch
 import torch.nn.functional as F
 import torch_xla
 
-from ai_edge_torch.hlfb import StableHLOCompositeBuilder
-
 
 def _export_stablehlo_mlir(model, args):
   ep = torch.export.export(model, args)
@@ -80,7 +79,9 @@ class TestStableHLOCompositeBuilder(unittest.TestCase):
         super().__init__()
 
       def log_softmax(self, x: torch.Tensor, dim: int):
-        builder = StableHLOCompositeBuilder(name="test.log_softmax", attr={"dim": dim})
+        builder = StableHLOCompositeBuilder(
+            name="test.log_softmax", attr={"dim": dim}
+        )
         x = builder.mark_inputs(x)
         y = torch.nn.functional.log_softmax(x, dim=dim)
         y = builder.mark_outputs(y)
@@ -126,7 +127,8 @@ class TestStableHLOCompositeBuilder(unittest.TestCase):
     self.assertEqual(mlir.count('stablehlo.composite "test.log_softmax"'), 1)
     self.assertEqual(
         mlir.count(
-            'composite_attributes = {dim = 0 : i64, source = "torch.nn", version = 1.000000e+00 : f32}'
+            'composite_attributes = {dim = 0 : i64, source = "torch.nn",'
+            " version = 1.000000e+00 : f32}"
         ),
         1,
     )
@@ -236,8 +238,12 @@ class TestStableHLOCompositeBuilder(unittest.TestCase):
     self.assertEqual(
         mlir.count('stablehlo.composite "test.scaled_dot_product_attention"'), 2
     )
-    self.assertEqual(mlir.count("composite_attributes = {include_captanh = true}"), 1)
-    self.assertEqual(mlir.count("composite_attributes = {include_captanh = false}"), 1)
+    self.assertEqual(
+        mlir.count("composite_attributes = {include_captanh = true}"), 1
+    )
+    self.assertEqual(
+        mlir.count("composite_attributes = {include_captanh = false}"), 1
+    )
 
   def test_build_composite_with_multiple_inputs_outputs(self):
     class SampleModel(torch.nn.Module):

ai_edge_torch/model.py

@@ -21,12 +21,11 @@ from __future__ import annotations
 
 import abc
 
+from ai_edge_torch.convert import conversion_utils as cutils
 import numpy as np
 import numpy.typing as npt
 import tensorflow as tf
 
-from ai_edge_torch.convert import conversion_utils as cutils
-
 
 class Model(abc.ABC):
   """Represents and edge model."""
@@ -84,7 +83,8 @@ class TfLiteModel(Model):
     signature_list = interpreter.get_signature_list()
     if signature_name not in signature_list:
       raise ValueError(
-          f"Invalid signature name provided. Available signatures: {', '.join(signature_list.keys())}"
+          'Invalid signature name provided. Available signatures:'
+          f' {", ".join(signature_list.keys())}'
       )
 
     try:
@@ -92,14 +92,17 @@ class TfLiteModel(Model):
     except ValueError as exception:
       if 'Invalid signature_key provided.' in str(exception):
         raise ValueError(
-            f'Invalid signature key provided. Available signatures: {list(signature_list.keys())}'
+            'Invalid signature key provided. Available signatures:'
+            f' {list(signature_list.keys())}'
         )
       else:
         raise exception
 
     if len(signature_list[signature_name]['inputs']) != len(args) + len(kwargs):
       raise ValueError(
-          f"The model requires {len(signature_list[signature_name]['inputs'])} arguments but {len(args)} was provided."
+          'The model requires'
+          f' {len(signature_list[signature_name]["inputs"])} arguments but'
+          f' {len(args)} was provided.'
       )
 
     # Gather the input dictionary based on the signature.

ai_edge_torch/quantize/pt2e_quantizer.py

@@ -19,6 +19,12 @@ import copy
 import functools
 from typing import Any, Callable, Dict, List, Optional, Set
 
+from ai_edge_torch.quantize.pt2e_quantizer_utils import _convert_scalars_to_attrs  # NOQA
+from ai_edge_torch.quantize.pt2e_quantizer_utils import OP_TO_ANNOTATOR
+from ai_edge_torch.quantize.pt2e_quantizer_utils import OperatorConfig
+from ai_edge_torch.quantize.pt2e_quantizer_utils import OperatorPatternType
+from ai_edge_torch.quantize.pt2e_quantizer_utils import propagate_annotation
+from ai_edge_torch.quantize.pt2e_quantizer_utils import QuantizationConfig
 import torch
 from torch.ao.quantization.fake_quantize import FusedMovingAvgObsFakeQuantize
 from torch.ao.quantization.observer import HistogramObserver
@@ -34,20 +40,15 @@ from torch.ao.quantization.quantizer import Quantizer
 from torch.fx import Node
 import torch.nn.functional as F
 
-from ai_edge_torch.quantize.pt2e_quantizer_utils import _convert_scalars_to_attrs  # NOQA
-from ai_edge_torch.quantize.pt2e_quantizer_utils import OP_TO_ANNOTATOR
-from ai_edge_torch.quantize.pt2e_quantizer_utils import OperatorConfig
-from ai_edge_torch.quantize.pt2e_quantizer_utils import OperatorPatternType
-from ai_edge_torch.quantize.pt2e_quantizer_utils import propagate_annotation
-from ai_edge_torch.quantize.pt2e_quantizer_utils import QuantizationConfig
-
 __all__ = [
     "PT2EQuantizer",
     "get_symmetric_quantization_config",
 ]
 
 
-def _supported_symmetric_quantized_operators() -> Dict[str, List[OperatorPatternType]]:
+def _supported_symmetric_quantized_operators() -> (
+    Dict[str, List[OperatorPatternType]]
+):
   supported_operators: Dict[str, List[OperatorPatternType]] = {
       # Both conv and linear should be able to handle relu + hardtanh fusion since
       # those are clamp ops
@@ -92,7 +93,9 @@ def get_symmetric_quantization_config(
 ):
   if is_qat:
     if is_dynamic:
-      raise NotImplementedError("dynamic quantization for qat is not yet implemented.")
+      raise NotImplementedError(
+          "dynamic quantization for qat is not yet implemented."
+      )
     act_observer_or_fake_quant_ctr = FusedMovingAvgObsFakeQuantize
   else:
     if is_dynamic:
@@ -106,12 +109,18 @@ def get_symmetric_quantization_config(
       quant_max=127,
       qscheme=torch.per_tensor_affine,
       is_dynamic=is_dynamic,
-      observer_or_fake_quant_ctr=act_observer_or_fake_quant_ctr.with_args(eps=2**-12),
+      observer_or_fake_quant_ctr=act_observer_or_fake_quant_ctr.with_args(
+          eps=2**-12
+      ),
   )
   qscheme = (
-      torch.per_channel_symmetric if is_per_channel else torch.per_tensor_symmetric
+      torch.per_channel_symmetric
+      if is_per_channel
+      else torch.per_tensor_symmetric
+  )
+  weight_observer_or_fake_quant_ctr: _ObserverOrFakeQuantizeConstructor = (
+      MinMaxObserver
   )
-  weight_observer_or_fake_quant_ctr: _ObserverOrFakeQuantizeConstructor = MinMaxObserver
   if is_qat:
     weight_observer_or_fake_quant_ctr = FusedMovingAvgObsFakeQuantize
   elif is_per_channel:
@@ -197,7 +206,9 @@ def _get_module_name_filter(module_name: str):
     # }
     # get_attr nodes doesn't have nn_module_stack?
     nn_module_stack = n.meta.get("nn_module_stack", {})
-    names = [n[len("L__self___") :].replace("_", ".") for n in nn_module_stack.keys()]
+    names = [
+        n[len("L__self___") :].replace("_", ".") for n in nn_module_stack.keys()
+    ]
     return module_name in names
 
   return module_name_filter
@@ -232,7 +243,9 @@ def _get_not_module_type_or_name_filter(
     tp_list: List[Callable], module_name_list: List[str]
 ) -> Callable[[Node], bool]:
   module_type_filters = [_get_module_type_filter(tp) for tp in tp_list]
-  module_name_list_filters = [_get_module_name_filter(m) for m in module_name_list]
+  module_name_list_filters = [
+      _get_module_name_filter(m) for m in module_name_list
+  ]
 
   def not_module_type_or_name_filter(n: Node) -> bool:
     return not any(f(n) for f in module_type_filters + module_name_list_filters)
@@ -307,7 +320,9 @@ class PT2EQuantizer(Quantizer):
         return ops
     return []
 
-  def set_global(self, quantization_config: QuantizationConfig) -> PT2EQuantizer:
+  def set_global(
+      self, quantization_config: QuantizationConfig
+  ) -> PT2EQuantizer:
     self.global_config = quantization_config
     return self
 

ai_edge_torch/quantize/pt2e_quantizer_utils.py

@@ -154,7 +154,9 @@ def get_weight_qspec(quantization_config: Optional[QuantizationConfig]):
       torch.per_tensor_symmetric,
       torch.per_channel_symmetric,
   ]:
-    raise ValueError(f"Unsupported quantization_spec {quantization_spec} for weight")
+    raise ValueError(
+        f"Unsupported quantization_spec {quantization_spec} for weight"
+    )
   return quantization_spec
 
 
@@ -193,7 +195,10 @@ def _annotate_linear(
   weight_qspec = get_weight_qspec(quantization_config)
   bias_qspec = get_bias_qspec(quantization_config)
   for node in gm.graph.nodes:
-    if node.op != "call_function" or node.target != torch.ops.aten.linear.default:
+    if (
+        node.op != "call_function"
+        or node.target != torch.ops.aten.linear.default
+    ):
       continue
     if filter_fn and not filter_fn(node):
       continue
@@ -417,7 +422,9 @@ def _annotate_conv_bn(
   Find conv + batchnorm parititions
   Note: This is only used for QAT. In PTQ, batchnorm should already be fused into the conv.
   """
-  return _do_annotate_conv_bn(gm, quantization_config, filter_fn, has_relu=False)
+  return _do_annotate_conv_bn(
+      gm, quantization_config, filter_fn, has_relu=False
+  )
 
 
 @register_annotator("conv_bn_relu")
@@ -486,7 +493,9 @@ def _do_annotate_conv_bn(
   # Match against all conv dimensions and cuda variants
   for (conv_fn, example_inputs), is_cuda, relu_is_inplace in combinations:
     pattern = get_pattern(conv_fn, relu_is_inplace)
-    pattern = _get_aten_graph_module_for_pattern(pattern, example_inputs, is_cuda)
+    pattern = _get_aten_graph_module_for_pattern(
+        pattern, example_inputs, is_cuda
+    )
     pattern.graph.eliminate_dead_code()
     pattern.recompile()
     matcher = SubgraphMatcherWithNameNodeMap(pattern, ignore_literals=True)
@@ -676,7 +685,9 @@ def _annotate_adaptive_avg_pool2d(
         and pool_node.target != torch.ops.aten.mean.dim
         and pool_node.target != torch.ops.aten.as_strided_.default
     ):
-      raise ValueError(f"{pool_node} is not an aten adaptive_avg_pool2d operator")
+      raise ValueError(
+          f"{pool_node} is not an aten adaptive_avg_pool2d operator"
+      )
 
     if _is_annotated([pool_node]):
       continue
@@ -741,7 +752,8 @@ def _annotate_fixed_qparams(
       continue
 
     node.meta["quantization_annotation"] = QuantizationAnnotation(
-        output_qspec=get_fixed_qparams_qspec(quantization_config), _annotated=True
+        output_qspec=get_fixed_qparams_qspec(quantization_config),
+        _annotated=True,
     )
     _mark_nodes_as_annotated(partition)
     annotated_partitions.append(partition)
@@ -885,7 +897,9 @@ def _annotate_mul(
     filter_fn: Optional[Callable[[Node], bool]] = None,
 ) -> Optional[List[List[Node]]]:
   mul_partitions = get_source_partitions(
-      gm.graph, ["mul", "mul_", operator.mul, torch.mul, operator.imul], filter_fn
+      gm.graph,
+      ["mul", "mul_", operator.mul, torch.mul, operator.imul],
+      filter_fn,
   )
   mul_partitions = list(itertools.chain(*mul_partitions.values()))
   annotated_partitions = []
@@ -932,8 +946,9 @@ def _annotate_cat(
 
     if cat_node.target != torch.ops.aten.cat.default:
       raise Exception(
-          f"Expected cat node: torch.ops.aten.cat.default, but found {cat_node.target}"
-          " please check if you are calling the correct capture API"
+          "Expected cat node: torch.ops.aten.cat.default, but found"
+          f" {cat_node.target} please check if you are calling the correct"
+          " capture API"
       )
 
     annotated_partitions.append(cat_partition.nodes)
@@ -987,7 +1002,9 @@ def propagate_annotation(model: torch.fx.GraphModule) -> None:
     if not isinstance(prev_node, Node):
       continue
 
-    quantization_annotation = prev_node.meta.get("quantization_annotation", None)
+    quantization_annotation = prev_node.meta.get(
+        "quantization_annotation", None
+    )
     if not quantization_annotation:
       continue
 
@@ -1014,7 +1031,9 @@ def propagate_annotation(model: torch.fx.GraphModule) -> None:
 
 
 # TODO: make the list of ops customizable
-def _convert_scalars_to_attrs(model: torch.fx.GraphModule) -> torch.fx.GraphModule:
+def _convert_scalars_to_attrs(
+    model: torch.fx.GraphModule,
+) -> torch.fx.GraphModule:
   for n in model.graph.nodes:
     if n.op != "call_function" or n.target not in [
         torch.ops.aten.add.Tensor,

ai_edge_torch/quantize/quant_config.py

@@ -76,6 +76,10 @@ class QuantConfig:
     elif generative_recipe is not None:
       generative_recipe.verify()
       object.__setattr__(self, 'generative_recipe', generative_recipe)
-      object.__setattr__(self, '_quantizer_mode', self._QuantizerMode.AI_EDGE_QUANTIZER)
+      object.__setattr__(
+          self, '_quantizer_mode', self._QuantizerMode.AI_EDGE_QUANTIZER
+      )
     else:
-      raise ValueError('Either pt2e_quantizer or generative_recipe must be set.')
+      raise ValueError(
+          'Either pt2e_quantizer or generative_recipe must be set.'
+      )