ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl → 0.3.0.dev20240926__py3-none-any.whl

ai_edge_torch/odml_torch/tf_integration.py

@@ -0,0 +1,194 @@
+# 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.
+# ==============================================================================
+"""APIs to convert lowered MLIR from PyTorch to TensorFlow and TFLite artifacts."""
+
+import re
+import tempfile
+
+import tensorflow as tf
+import torch
+
+from tensorflow.compiler.tf2xla.python import xla as tfxla
+
+from . import export
+from . import export_utils
+
+
+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 _mangle_tf_root_scope_name(name):
+  r"""Build the mangled name for tf.Variable.
+
+  TF has more restricted constrain on the variable names at root scope. Root
+  scope name constrain: [A-Za-z0-9.][A-Za-z0-9_.\\-/]* Non-root scope name
+  constrain: [A-Za-z0-9_.\\-/]*
+  https://github.com/tensorflow/tensorflow/blob/51b601fa6bb7e801c0b6ae73c25580e40a8b5745/tensorflow/python/framework/ops.py#L3301-L3302
+  The state_dict key doesn't have such constrain, the name need to be mangled
+  when a root-scoped TF variable is created.
+
+  FX Graph Node may contain characters other than [A-Za-z0-9_.\\-/], replace
+  offending characters with '_'.
+
+  Args:
+    name: the tensor name to be mangled.
+
+  Returns:
+    Mangled name in str.
+  """
+  if name[0] in "._\\-/":
+    name = "k" + name
+  name = re.sub(r"[^^\w\-/\\]+", "_", name)
+  return name
+
+
+def _build_tf_state_dict(
+    lowered: export.MlirLowered,
+) -> dict[str, tf.Variable]:
+  """Build a dictionary of tf.Variable from the state_dict in lowered."""
+  tf_state_dict = {}
+  for sig in lowered.input_signature:
+    if sig.input_spec.is_parameter:
+      name = sig.input_spec.name
+      tf_state_dict[name] = tf.Variable(
+          lowered.state_dict[name].detach().numpy(),
+          trainable=False,
+          name=_mangle_tf_root_scope_name(name),
+      )
+  return tf_state_dict
+
+
+def _extract_call_args(
+    lowered: export.MlirLowered,
+    args,
+    tf_state_dict: dict[str, tf.Variable],
+):
+  """Extract the flattened inputs to built tf.function."""
+  call_args = []
+  for sig in lowered.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(lowered, tf_state_dict):
+  """Build tf.function from lowered and tf_state_dict."""
+
+  def inner(*args):
+    t_outs = [torch_dtype_to_tf(sig.dtype) for sig in lowered.output_signature]
+    s_outs = [_get_shape_with_dynamic(sig) for sig in lowered.output_signature]
+    call_args = _extract_call_args(lowered, 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=lowered.module_bytecode,
+    )
+
+  return inner
+
+
+def _make_input_signatures(
+    lowered: export.MlirLowered,
+) -> list[tf.TensorSpec]:
+  """Build the input signatures in tf.TensorSpec for building tf.function."""
+  user_input_signature = sorted(
+      [sig for sig in lowered.input_signature if sig.input_spec.is_user_input],
+      key=lambda sig: sig.input_spec.i,
+  )
+  tf_signatures = []
+
+  for sig in user_input_signature:
+    shape = _get_shape_with_dynamic(sig)
+    tf_signatures.append(
+        tf.TensorSpec(
+            shape=shape,
+            dtype=torch_dtype_to_tf(sig.dtype),
+            name=f"args_{sig.input_spec.i}",
+        )
+    )
+  return tf_signatures
+
+
+def mlir_to_tf_function(lowered: export.MlirLowered):
+  """Convert the MLIR lowered to a executable tf.function."""
+  tf_state_dict = _build_tf_state_dict(lowered)
+  return tf.function(
+      _wrap_as_tf_func(lowered, tf_state_dict),
+      input_signature=_make_input_signatures(lowered),
+  )
+
+
+def mlir_to_flatbuffer(lowered: export.MlirLowered):
+  """Convert the MLIR lowered to a TFLite flatbuffer binary."""
+  tf_state_dict = _build_tf_state_dict(lowered)
+  signature_names = [tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
+  tf_signatures = [_make_input_signatures(lowered)]
+  tf_functions = [_wrap_as_tf_func(lowered, tf_state_dict)]
+
+  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 = list(tf_state_dict.values())
+
+  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_name in enumerate(signature_names)
+        },
+    )
+
+    converter = tf.lite.TFLiteConverter.from_saved_model(temp_dir_path)
+    tflite_model = converter.convert()
+
+  return tflite_model

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:
@@ -179,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:
@@ -197,7 +209,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
@@ -205,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:
@@ -232,7 +250,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 +327,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
 
@@ -323,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
@@ -333,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

ai_edge_torch/quantize/pt2e_quantizer_utils.py

@@ -31,7 +31,7 @@ from torch.ao.quantization.quantizer import SharedQuantizationSpec
 from torch.ao.quantization.quantizer.utils import _annotate_input_qspec_map
 from torch.ao.quantization.quantizer.utils import _annotate_output_qspec
 from torch.fx import Node
-from torch.fx.passes.utils.matcher_with_name_node_map_utils import SubgraphMatcherWithNameNodeMap  # NOQA
+from torch.fx.passes.utils.matcher_with_name_node_map_utils import SubgraphMatcherWithNameNodeMap
 from torch.fx.passes.utils.source_matcher_utils import get_source_partitions
 import torch.nn.functional as F
 
@@ -95,9 +95,10 @@ class OperatorConfig(NamedTuple):
 
 
 def _is_annotated(nodes: List[Node]):
-  """
-  Given a list of nodes (that represents an operator pattern),
-  check if any of the node is annotated, return True if any of the node
+  """Checks if a list of nodes is annotated.
+
+  Given a list of nodes (that represents an operator pattern), check if any of
+  the node is annotated, return True if any of the node
   is annotated, otherwise return False
   """
   annotated = False
@@ -154,7 +155,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 +196,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
@@ -413,11 +419,13 @@ def _annotate_conv_bn(
     quantization_config: Optional[QuantizationConfig],
     filter_fn: Optional[Callable[[Node], bool]] = None,
 ) -> Optional[List[List[Node]]]:
+  """Find conv + batchnorm parititions Note: This is only used for QAT.
+
+  In PTQ, batchnorm should already be fused into the conv.
   """
-  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")
@@ -426,9 +434,9 @@ def _annotate_conv_bn_relu(
     quantization_config: Optional[QuantizationConfig],
     filter_fn: Optional[Callable[[Node], bool]] = None,
 ) -> Optional[List[List[Node]]]:
-  """
-  Find conv + batchnorm + relu parititions
-  Note: This is only used for QAT. In PTQ, batchnorm should already be fused into the conv.
+  """Find conv + batchnorm + relu 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=True)
 
@@ -439,8 +447,8 @@ def _do_annotate_conv_bn(
     filter_fn: Optional[Callable[[Node], bool]],
     has_relu: bool,
 ) -> List[List[Node]]:
-  """
-  Given a function that takes in a `conv_fn` and returns a conv-bn[-relu] pattern,
+  """Given a function that takes in a `conv_fn` and returns a conv-bn[-relu] pattern,
+
   return a list of annotated partitions.
 
   The output of the pattern must include a dictionary from string name to node
@@ -486,7 +494,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 +686,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 +753,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 +898,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 +947,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 +1003,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 +1032,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

@@ -13,27 +13,27 @@
 # limitations under the License.
 # ==============================================================================
 
-from dataclasses import dataclass
+import dataclasses
 import enum
 from typing import Optional
 
-from ai_edge_torch.generative.quantize import quant_attrs
 from ai_edge_torch.generative.quantize import quant_recipe
 from ai_edge_torch.quantize import pt2e_quantizer as pt2eq
 
 
-@dataclass(frozen=True)
+@dataclasses.dataclass(frozen=True)
 class QuantConfig:
-  """
+  """Encapsulates a quantization configuration.
+
   Encapsulates all different quantization methods and schemes available for
   models converted with ai_edge_torch.
 
-  Args:
+  Attributes:
     pt2e_quantizer: The instance of PT2EQuantizer used to quantize the model
       with PT2E quantization. This method of quantization is not applicable to
       models created with the Edge Generative API.
-    generative_recipe: Quantization recipe to be applied on a model created
-      with the Edge Generative API.
+    generative_recipe: Quantization recipe to be applied on a model created with
+      the Edge Generative API.
   """
 
   pt2e_quantizer: pt2eq.PT2EQuantizer = None
@@ -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.'
+      )

ai_edge_torch/testing/model_coverage/model_coverage.py

@@ -13,26 +13,33 @@
 # limitations under the License.
 # ==============================================================================
 
-"""Utility Functions to test TFLite models exported from PyTorch"""
+"""Contains utility functions to test TFLite models exported from PyTorch."""
 
 from collections.abc import Callable
 
+from ai_edge_torch import model
 import numpy as np
 import torch
 from torch.utils import _pytree as pytree
 
-from ai_edge_torch.model import Model
-
 
 # Utility to flatten the order to make it deterministic.
 # Ordering is done in left-to-right depth-first tree traversal.
 def _flatten(data):
-  out, spec = pytree.tree_flatten(data)
+  out, _ = pytree.tree_flatten(data)
   return out
 
 
 # Convert a Torch Tensor to a numpy array
 def _torch_tensors_to_np(*argv):
+  """Converts a Torch Tensor to a numpy array.
+
+  Args:
+    *argv: A list of torch.tensor or a single torch.tensor.
+
+  Returns:
+    A list of numpy array or a single numpy array.
+  """
   if len(argv) > 1:
     data = list(argv)
   else:
@@ -58,7 +65,7 @@ def _torch_tensors_to_np(*argv):
 
 
 def compare_tflite_torch(
-    edge_model: Model,
+    edge_model: model.Model,
     torch_eval_func: Callable,
     args=None,
     kwargs=None,
@@ -69,15 +76,17 @@ def compare_tflite_torch(
     rtol: float = 1e-5
 ):
   """Compares torch models and TFLite models.
+
   Args:
     edge_model: Serialized ai_edge_torch.model.Model object.
     torch_eval_func: Callable function to evaluate torch model.
-    args: torch.tensor array or a callable to generate a torch.tensor array
-      with random data, to pass into models during inference. (default None).
+    args: torch.tensor array or a callable to generate a torch.tensor array with
+      random data, to pass into models during inference. (default None).
     kwargs: dict of str to torch.tensor, or a callable to generate such.
-    num_valid_inputs: Defines the number of times the random inputs will be generated (if a callable is provided for input_data).
-    signature_name: If provided, specifies the name for the signature of the edge_model to run.
-      Calls the default signature if not provided.
+    num_valid_inputs: Defines the number of times the random inputs will be
+      generated (if a callable is provided for input_data).
+    signature_name: If provided, specifies the name for the signature of the
+      edge_model to run. Calls the default signature if not provided.
     atol: Absolute tolerance (see `numpy.allclose`)
     rtol: Relative tolerance (see `numpy.allclose`)
   """
@@ -94,7 +103,9 @@ def compare_tflite_torch(
       )
       for _ in range(num_valid_inputs)
   ]
-  torch_outputs = [torch_eval_func(*args, **kwargs) for args, kwargs in torch_inputs]
+  torch_outputs = [
+      torch_eval_func(*args, **kwargs) for args, kwargs in torch_inputs
+  ]
   np_inputs = [
       (_torch_tensors_to_np(args), _torch_tensors_to_np(kwargs))
       for args, kwargs in torch_inputs
@@ -110,12 +121,13 @@ def compare_tflite_torch(
     if signature_name is None:
       return _flatten(edge_model(*args, **kwargs))
     else:
-      return _flatten(edge_model(*args, **kwargs, signature_name=signature_name))
+      return _flatten(
+          edge_model(*args, **kwargs, signature_name=signature_name)
+      )
 
   for idx, np_input in enumerate(np_inputs):
     output = get_edge_output(np_input)
     golden_output = np_outputs[idx]
-
     is_output_len_eq = len(golden_output) == len(output)
 
     output = [v.astype(np.float32) for v in output]
@@ -123,9 +135,10 @@ def compare_tflite_torch(
 
     # Append the results of each invoke to a function-global variable
     # used to store the comparison final results
-    is_equal = is_output_len_eq and all(
-        [equal_fn(out, golden_out) for out, golden_out in zip(output, golden_output)]
-    )
+    is_equal = is_output_len_eq and all([
+        equal_fn(out, golden_out)
+        for out, golden_out in zip(output, golden_output)
+    ])
     if not is_equal:
       return False