ai-edge-torch-nightly 0.2.0.dev20240806__py3-none-any.whl → 0.2.0.dev20240808__py3-none-any.whl

ai_edge_torch/lowertools/common_utils.py

@@ -0,0 +1,89 @@
+# 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 logging
+
+from ai_edge_torch._convert import signature as signature_module
+import tensorflow as tf
+import torch
+
+
+def _torch_to_tf_variable(torch_tensor: torch.Tensor):
+  if not torch_tensor.is_contiguous():
+    torch_tensor = torch_tensor.contiguous()
+
+  try:
+    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."
+    )
+    nparray = torch_tensor.cpu().detach().numpy()
+    tf_tensor = tf.convert_to_tensor(nparray)
+
+  return tf.Variable(tf_tensor, trainable=False)
+
+
+def _get_states(
+    exported_programs: list[torch.export.ExportedProgram],
+    signatures: list[signature_module.Signature],
+):
+  for exported_program, signature in zip(exported_programs, signatures):
+    args, _ = exported_program.example_inputs
+    # Calling this to get **all** the state including model buffers.
+    _flat_input_args = exported_program._graph_module_flat_inputs(args, {})
+    for tensor, input_spec in zip(
+        _flat_input_args, exported_program.graph_signature.input_specs
+    ):
+      # 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
+      ):
+        continue
+      yield signature, tensor, input_spec
+
+
+def _tensor_unique_id(tensor: torch.Tensor):
+  return (
+      str(tensor.device),
+      tensor.shape,
+      tensor.stride(),
+      tensor.untyped_storage().data_ptr(),
+  )
+
+
+def gather_state_dict(
+    exported_programs: list[torch.export.ExportedProgram],
+    signatures: list[signature_module.Signature],
+):
+  deduped_tensor_map = {}
+
+  for _, tensor, _ in _get_states(exported_programs, signatures):
+    unique_id = _tensor_unique_id(tensor)
+    deduped_tensor_map[unique_id] = _torch_to_tf_variable(tensor)
+
+  state_dict = {}
+  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
+    ]
+
+  return state_dict, list(deduped_tensor_map.values())

ai_edge_torch/lowertools/odml_torch_utils.py

@@ -0,0 +1,211 @@
+# 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 dataclasses
+import tempfile
+from typing import Any, Dict, List, Optional, Tuple
+
+from ai_edge_torch import odml_torch
+from ai_edge_torch._convert import conversion_utils
+from ai_edge_torch._convert import signature as signature_module
+from ai_edge_torch.lowertools import common_utils
+from ai_edge_torch.odml_torch import export
+from ai_edge_torch.odml_torch import export_utils
+from ai_edge_torch.quantize import quant_config as qcfg
+import tensorflow as tf
+import torch
+
+from tensorflow.compiler.tf2xla.python import xla as tfxla
+
+MlirBundle = odml_torch.export.MlirLowered
+
+
+@dataclasses.dataclass
+class MergedBundle:
+  """A bundle of MlirLowered that has been merged."""
+
+  bundles: list[odml_torch.export.MlirLowered]
+  deduped_tf_vars: list[tf.Variable]
+
+
+def torch_dtype_to_tf(dtype):
+  return {
+      torch.double: tf.float64,
+      torch.float32: tf.float32,
+      torch.half: tf.float16,
+      torch.long: tf.int64,
+      torch.int32: tf.int32,
+      torch.int16: tf.int16,
+      torch.bool: tf.bool,
+  }.get(dtype)
+
+
+def _get_shape_with_dynamic(signature: export.VariableSignature):
+  return [
+      None if export_utils.is_torch_dynamic(s) else s for s in signature.shape
+  ]
+
+
+def _extract_call_args(
+    bundle: export.MlirLowered,
+    args: Tuple[Any],
+    tf_state_dict: Dict[str, tf.Variable],
+):
+  call_args = []
+  for sig in bundle.input_signature:
+    if sig.input_spec.is_user_input:
+      call_args.append(args[sig.input_spec.i])
+    elif sig.input_spec.is_parameter:
+      name = sig.input_spec.name
+      call_args.append(tf_state_dict[name])
+  return call_args
+
+
+def _wrap_as_tf_func(bundle, tf_state_dict):
+  def inner(*args):
+    t_outs = [torch_dtype_to_tf(sig.dtype) for sig in bundle.output_signature]
+    s_outs = [_get_shape_with_dynamic(sig) for sig in bundle.output_signature]
+    call_args = _extract_call_args(bundle, args, tf_state_dict)
+    return tfxla.call_module(
+        tuple(call_args),
+        version=5,
+        Tout=t_outs,  # dtype information
+        Sout=s_outs,  # Shape information
+        function_list=[],
+        module=bundle.module_bytecode,
+    )
+
+  return inner
+
+
+def _make_tf_signature(
+    input_signature: list[export.VariableSignature],
+    signature: signature_module.Signature,
+) -> List[tf.TensorSpec]:
+  input_names = signature.flat_arg_names
+  user_input_signature = sorted(
+      [sig for sig in input_signature if sig.input_spec.is_user_input],
+      key=lambda sig: sig.input_spec.i,
+  )
+  tf_signature = []
+
+  for sig in user_input_signature:
+    shape = _get_shape_with_dynamic(sig)
+    tf_signature.append(
+        tf.TensorSpec(
+            shape=shape,
+            dtype=torch_dtype_to_tf(sig.dtype),
+            name=input_names[sig.input_spec.i],
+        )
+    )
+  return tf_signature
+
+
+def merged_bundle_to_tfl_model(
+    merged_bundle: MergedBundle,
+    signatures: list[signature_module.Signature],
+    *,
+    quant_config: Optional[qcfg.QuantConfig] = None,
+    _tfl_converter_flags: dict = {},
+):
+  tf_state_dict = merged_bundle.bundles[0].state_dict
+
+  tf_signatures = [
+      _make_tf_signature(bundle.input_signature, sig)
+      for bundle, sig in zip(merged_bundle.bundles, signatures)
+  ]
+  tf_functions = [
+      _wrap_as_tf_func(bundle, tf_state_dict)
+      for bundle in merged_bundle.bundles
+  ]
+
+  tf_module = tf.Module()
+  tf_module.f = []
+
+  for tf_sig, func in zip(tf_signatures, tf_functions):
+    tf_module.f.append(
+        tf.function(
+            func,
+            input_signature=tf_sig,
+        )
+    )
+
+  tf_module._variables = merged_bundle.deduped_tf_vars
+
+  tf_concrete_funcs = [
+      func.get_concrete_function(*tf_sig)
+      for func, tf_sig in zip(tf_module.f, tf_signatures)
+  ]
+
+  # We need to temporarily save since TFLite's from_concrete_functions does not
+  # allow providing names for each of the concrete functions.
+  with tempfile.TemporaryDirectory() as temp_dir_path:
+    tf.saved_model.save(
+        tf_module,
+        temp_dir_path,
+        signatures={
+            sig.name: tf_concrete_funcs[idx]
+            for idx, sig in enumerate(signatures)
+        },
+    )
+
+    converter = tf.lite.TFLiteConverter.from_saved_model(temp_dir_path)
+    converter._experimental_enable_composite_direct_lowering = True
+
+    conversion_utils.apply_tfl_converter_flags(converter, _tfl_converter_flags)
+
+    tflite_model = converter.convert()
+
+  return tflite_model
+
+
+def exported_program_to_mlir_text(
+    exported_program: torch.export.ExportedProgram,
+) -> str:
+  """Converts a ExportedProgram to a MLIR text."""
+  return odml_torch.export.exported_program_to_mlir(exported_program).get_text(
+      enable_debug_info=True
+  )
+
+
+def exported_program_to_mlir(
+    exported_program: torch.export.ExportedProgram,
+    sample_args: tuple[torch.Tensor],
+) -> export.MlirLowered:
+  """Converts a ExportedProgram to a MlirLowered."""
+  return odml_torch.export.exported_program_to_mlir(exported_program)
+
+
+def merge_mlir_bundles(
+    bundles: list[export.MlirLowered],
+    signatures: list[signature_module.Signature],
+    exported_programs: list[torch.export.ExportedProgram],
+) -> MergedBundle:
+  """Merges a list of MlirLowered into one."""
+  state_dict, deduped_vars = common_utils.gather_state_dict(
+      exported_programs, signatures
+  )
+
+  merged_bundle = MergedBundle(
+      bundles=bundles.copy(), deduped_tf_vars=deduped_vars
+  )
+  for bundle, signature in zip(merged_bundle.bundles, signatures):
+    bundle.state_dict = state_dict
+
+    for var_sig in bundle.input_signature:
+      if var_sig.input_spec.is_parameter:
+        var_sig.input_spec.name = signature.name + "_" + var_sig.input_spec.name
+
+  return merged_bundle

ai_edge_torch/lowertools/torch_xla_utils.py

@@ -0,0 +1,273 @@
+# 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 dataclasses
+from dataclasses import dataclass
+import gc
+import itertools
+import logging
+import tempfile
+from typing import Any, Dict, Optional, Tuple, Union
+
+from ai_edge_torch import model
+from ai_edge_torch._convert import conversion_utils
+from ai_edge_torch._convert import signature as signature_module
+from ai_edge_torch.generative.quantize.ai_edge_quantizer_glue import translate_recipe  # NOQA
+from ai_edge_torch.lowertools import common_utils
+from ai_edge_torch.quantize import quant_config as qcfg
+import torch
+from torch_xla import stablehlo
+
+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
+except ImportError:
+  logging.error(
+      "This module needs tensorflow with xla support.\n"
+      "Please install tensorflow with `pip install tf-nightly`.\n"
+  )
+  raise
+
+MlirBundle = stablehlo.StableHLOModelBundle
+
+
+@dataclasses.dataclass
+class MergedBundle:
+
+  bundle: stablehlo.StableHLOModelBundle
+  deduped_tf_vars: list[tf.Variable]
+
+
+def exported_program_to_mlir(
+    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
+  # in the returned bundle is going to be empty.
+  return stablehlo.exported_program_to_stablehlo(
+      exported_program,
+      stablehlo.StableHLOExportOptions(
+          override_tracing_arguments=sample_args, export_weights=False
+      ),
+  )._bundle
+
+
+def merge_mlir_bundles(
+    bundles: list[stablehlo.StableHLOModelBundle],
+    signatures: list[signature_module.Signature],
+    exported_programs: list[torch.export.ExportedProgram],
+) -> stablehlo.StableHLOGraphModule:
+  state_dict, deduped_tf_vars = common_utils.gather_state_dict(
+      exported_programs, signatures
+  )
+
+  new_shlo_model_bundle = stablehlo.StableHLOModelBundle(
+      state_dict=state_dict, additional_constants=[], stablehlo_funcs=[]
+  )
+
+  for bundle, signature in zip(bundles, signatures):
+    const_offset = len(new_shlo_model_bundle.additional_constants)
+    for func in bundle.stablehlo_funcs:
+      func.meta.name = signature.name + "_" + func.meta.name
+      for loc in func.meta.input_locations:
+        if loc.type_ == stablehlo.VariableType.CONSTANT:
+          loc.position += const_offset
+        elif loc.type_ == stablehlo.VariableType.PARAMETER:
+          loc.name = signature.name + "_" + loc.name
+      new_shlo_model_bundle.stablehlo_funcs.append(func)
+    new_shlo_model_bundle.additional_constants.extend(
+        bundle.additional_constants
+    )
+  return MergedBundle(
+      bundle=new_shlo_model_bundle, deduped_tf_vars=deduped_tf_vars
+  )
+
+
+def _get_shape_with_dynamic(signature: stablehlo.VariableSignature):
+  shape = copy.copy(signature.shape)
+  for i in signature.dynamic_dims:
+    shape[i] = None
+  return shape
+
+
+def _wrap_as_tf_func(
+    func: stablehlo.StableHLOFunc, bundle: stablehlo.StableHLOModelBundle
+):
+  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
+    ]
+    call_args = stablehlo._extract_call_parameters(args, func.meta, bundle)
+    return tfxla.call_module(
+        tuple(call_args),
+        version=5,
+        Tout=type_info,
+        Sout=shape_info,
+        function_list=[],
+        module=func.bytecode,
+    )
+
+  return inner
+
+
+def _make_tf_function(
+    bundle: stablehlo.StableHLOModelBundle = None,
+):
+  bundle = bundle if bundle is None else bundle
+  return [_wrap_as_tf_func(func, bundle) for func in bundle.stablehlo_funcs]
+
+
+def _make_tf_signature(
+    meta: stablehlo.StableHLOFunctionMeta,
+    signature: signature_module.Signature,
+) -> list[tf.TensorSpec]:
+  input_names = signature.flat_arg_names
+  input_pos_to_spec = {
+      loc.position: spec
+      for loc, spec in itertools.chain(
+          zip(meta.input_locations, meta.input_signature), meta.unused_inputs
+      )
+      if loc.type_ == stablehlo.VariableType.INPUT_ARG
+  }
+  assert len(input_pos_to_spec) == len(input_names)
+
+  primitive_type_to_tf_type = {"int": "int32", "float": "float32"}
+  ret: list[tf.TensorSpec] = []
+  for i, name in enumerate(input_names):
+    spec = input_pos_to_spec[i]
+    shape = _get_shape_with_dynamic(spec)
+    ret.append(
+        tf.TensorSpec(
+            shape=shape,
+            dtype=primitive_type_to_tf_type[spec.dtype]
+            if spec.dtype in primitive_type_to_tf_type
+            else spec.dtype,
+            name=name,
+        )
+    )
+  return ret
+
+
+def exported_program_to_mlir_text(
+    exported_program: torch.export.ExportedProgram,
+) -> str:
+  """Converts a ExportedProgram to a MLIR text."""
+  return stablehlo.exported_program_to_stablehlo(
+      exported_program
+  ).get_stablehlo_text()
+
+
+def merged_bundle_to_tfl_model(
+    merged_bundle: MergedBundle,
+    signatures: list[signature_module.Signature],
+    *,
+    quant_config: Optional[qcfg.QuantConfig] = None,
+    _tfl_converter_flags: dict = {},
+) -> None:
+  """Converts a StableHLOGraphModule to a tflite model.
+
+  Args: shlo_bundle - model to export and save
+
+    signatures: List of signatures from which names of the signatures is
+    extracted.
+    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.
+  """
+
+  tf_module = tf.Module()
+
+  shlo_bundle = merged_bundle.bundle
+
+  shlo_bundle.additional_constants = [
+      tf.Variable(v, trainable=False) for v in shlo_bundle.additional_constants
+  ]
+  tf_signatures: list[list[tf.TensorSpec]] = list(
+      _make_tf_signature(func.meta, sig)
+      for func, sig in zip(shlo_bundle.stablehlo_funcs, signatures)
+  )
+
+  tf_functions = _make_tf_function(shlo_bundle)
+
+  tf_module.f = []
+  for tf_sig, func in zip(tf_signatures, tf_functions):
+    tf_module.f.append(
+        tf.function(
+            func,
+            input_signature=tf_sig,
+        )
+    )
+
+  tf_module._variables = (
+      merged_bundle.deduped_tf_vars + shlo_bundle.additional_constants
+  )
+  del shlo_bundle
+  gc.collect()
+
+  tf_concrete_funcs = [
+      func.get_concrete_function(*tf_sig)
+      for func, tf_sig in zip(tf_module.f, tf_signatures)
+  ]
+
+  # We need to temporarily save since TFLite's from_concrete_functions does not
+  # allow providing names for each of the concrete functions.
+  with tempfile.TemporaryDirectory() as temp_dir_path:
+    tf.saved_model.save(
+        tf_module,
+        temp_dir_path,
+        signatures={
+            sig.name: tf_concrete_funcs[idx]
+            for idx, sig in enumerate(signatures)
+        },
+    )
+    # Clean up intermediate memory early.
+    del tf_module
+    del tf_concrete_funcs
+    gc.collect()
+
+    converter = tf.lite.TFLiteConverter.from_saved_model(temp_dir_path)
+    converter._set_original_model_type(conversion_metadata_fb.ModelType.PYTORCH)
+    converter._experimental_enable_composite_direct_lowering = True
+
+    conversion_utils.set_tfl_converter_quant_flags(converter, quant_config)
+    if (
+        quant_config is not None
+        and quant_config._quantizer_mode
+        == quant_config._QuantizerMode.AI_EDGE_QUANTIZER
+    ):
+      translated_recipe = translate_recipe.translate_to_ai_edge_recipe(
+          quant_config.generative_recipe
+      )
+
+    conversion_utils.apply_tfl_converter_flags(converter, _tfl_converter_flags)
+
+    tflite_model = converter.convert()
+
+    if (
+        quant_config is not None
+        and quant_config._quantizer_mode
+        == quant_config._QuantizerMode.AI_EDGE_QUANTIZER
+    ):
+      tflite_model = translate_recipe.quantize_model(
+          tflite_model, translated_recipe
+      )
+
+  return tflite_model

ai_edge_torch/model.py

@@ -15,17 +15,18 @@
 
 """Represents an ai_edge_torch model.
 
-PyTorch models can be converted to this representation through `ai_edge_torch.convert`.
+PyTorch models can be converted to this representation through
+`ai_edge_torch.convert`.
 """
 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
 
+DEFAULT_SIGNATURE_NAME = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
+
 
 class Model(abc.ABC):
   """Represents and edge model."""
@@ -34,7 +35,7 @@ class Model(abc.ABC):
   def __call__(
       self,
       *args: npt.ArrayLike,
-      signature_name: str = cutils.DEFAULT_SIGNATURE_NAME,
+      signature_name: str = DEFAULT_SIGNATURE_NAME,
       **kwargs,
   ) -> npt.ArrayLike | tuple[npt.ArrayLike]:
     raise NotImplementedError()
@@ -66,18 +67,22 @@ class TfLiteModel(Model):
   def __call__(
       self,
       *args: npt.ArrayLike,
-      signature_name: str = cutils.DEFAULT_SIGNATURE_NAME,
+      signature_name: str = DEFAULT_SIGNATURE_NAME,
       **kwargs,
   ) -> npt.ArrayLike | tuple[npt.ArrayLike]:
     """Runs inference on the edge model using the provided arguments.
 
     Args:
       *args: The arguments to be passed to the model for inference.
-      **kwargs: The arguments with specific names to be passed to the model for inference.
-      signature_name: The name of the signature to be used for inference.
-        The default signature is used if not provided.
+      **kwargs: The arguments with specific names to be passed to the model for
+        inference.
+      signature_name: The name of the signature to be used for inference. The
+        default signature is used if not provided.
     """
-    interpreter = tf.lite.Interpreter(model_content=self._tflite_model)
+    interpreter = tf.lite.Interpreter(
+        model_content=self._tflite_model,
+        experimental_default_delegate_latest_features=True,
+    )
     interpreter.allocate_tensors()
 
     signature_list = interpreter.get_signature_list()

ai_edge_torch/quantize/pt2e_quantizer.py

@@ -188,15 +188,18 @@ def _get_supported_config_and_operators() -> List[OperatorConfig]:
 
 def _get_module_name_filter(module_name: str):
   """Get the module_name_filter function for a given module name, the filter accepts
+
   a node and checks if the node comes from a module that has certain module name
 
   For example:
-      node: linear_op = call_function[...](...)  # comes from a module with name blocks.sub.linear1
+      node: linear_op = call_function[...](...)  # comes from a module with name
+      blocks.sub.linear1
 
 
   >> module_name_filter = _get_module_name_filter("blocks.sub")
   >> print(module_name_filter(node))
-  True  # the node is from "blocks.sub" based on the fully qualified name "blocks.sub.linear1"
+  True  # the node is from "blocks.sub" based on the fully qualified name
+  "blocks.sub.linear1"
   """
 
   def module_name_filter(n: Node) -> bool:
@@ -216,15 +219,19 @@ def _get_module_name_filter(module_name: str):
 
 def _get_module_type_filter(tp: Callable):
   """Get the module_type_filter function for a given module type, the filter accepts
+
   a node and checks if the node comes from a module that has certain module type
 
   For example:
-      node: linear_op = call_function[...](...)  # comes from a module with type Block -> Sub -> Linear
+      node: linear_op = call_function[...](...)  # comes from a module with type
+      Block -> Sub -> Linear
 
 
-  >> module_type_filter = _get_module_type_filter(Sub)  # submodule with type `Sub`, under the `Block` submodule
+  >> module_type_filter = _get_module_type_filter(Sub)  # submodule with type
+  `Sub`, under the `Block` submodule
   >> print(module_type_filter(node))
-  True  # the node is from the submodule `Sub` (same for `Block` and `Linear` as well)
+  True  # the node is from the submodule `Sub` (same for `Block` and `Linear` as
+  well)
   """
 
   def module_type_filter(n: Node) -> bool:
@@ -338,8 +345,11 @@ class PT2EQuantizer(Quantizer):
       self, module_type: Callable, quantization_config: QuantizationConfig
   ):
     """Set quantization_config for a submodule with type: `module_type`, for example:
-    quantizer.set_module_name(Sub) or quantizer.set_module_name(nn.Linear), it will quantize all supported operator/operator
-    patterns in the submodule with this module type with the given `quantization_config`
+
+    quantizer.set_module_name(Sub) or quantizer.set_module_name(nn.Linear), it
+    will quantize all supported operator/operator
+    patterns in the submodule with this module type with the given
+    `quantization_config`
     """
     self.module_type_config[module_type] = quantization_config
     return self
@@ -348,8 +358,11 @@ class PT2EQuantizer(Quantizer):
       self, module_name: str, quantization_config: Optional[QuantizationConfig]
   ):
     """Set quantization_config for a submodule with name: `module_name`, for example:
-    quantizer.set_module_name("blocks.sub"), it will quantize all supported operator/operator
-    patterns in the submodule with this module name with the given `quantization_config`
+
+    quantizer.set_module_name("blocks.sub"), it will quantize all supported
+    operator/operator
+    patterns in the submodule with this module name with the given
+    `quantization_config`
     """
     assert (
         quantization_config is not None