fix: Linting and formatting updates

Author: gs-olive

py/torch_tensorrt/dynamo/_settings.py

@@ -10,10 +10,10 @@
     PASS_THROUGH_BUILD_FAILURES,
     PRECISION,
     TRUNCATE_LONG_AND_DOUBLE,
+    USE_FAST_PARTITIONER,
     USE_PYTHON_RUNTIME,
     VERSION_COMPATIBLE,
     WORKSPACE_SIZE,
-    USE_FAST_PARTITIONER,
 )
 
 

py/torch_tensorrt/dynamo/backend/backends.py

@@ -5,15 +5,7 @@
 import torch
 import torch._dynamo as td
 from torch._functorch.aot_autograd import aot_module_simplified, make_boxed_compiler
-from torch_tensorrt.dynamo import CompilationSettings
-from torch_tensorrt.dynamo.lowering._decompositions import (
-    get_decompositions,
-)
-from torch_tensorrt.dynamo.lowering._pre_aot_lowering import (
-    pre_aot_substitutions,
-)
-from torch_tensorrt.dynamo import partitioning
-from torch_tensorrt.dynamo.utils import parse_dynamo_kwargs
+from torch_tensorrt.dynamo import CompilationSettings, partitioning
 from torch_tensorrt.dynamo.conversion import (
     convert_module,
     repair_long_or_double_inputs,
@@ -138,7 +130,6 @@ def _compile_module(
     # Iterate over all components that can be accelerated
     # Generate the corresponding TRT Module for those
     for name, _ in partitioned_module.named_children():
-
         # Criteria for a module to be convertible to TRT
         if settings.use_fast_partitioner and "_run_on_acc" not in name:
             continue

py/torch_tensorrt/dynamo/compile.py

@@ -19,10 +19,10 @@
     PASS_THROUGH_BUILD_FAILURES,
     PRECISION,
     TRUNCATE_LONG_AND_DOUBLE,
+    USE_FAST_PARTITIONER,
     USE_PYTHON_RUNTIME,
     VERSION_COMPATIBLE,
     WORKSPACE_SIZE,
-    USE_FAST_PARTITIONER,
 )
 from torch_tensorrt.dynamo.backend.backends import _compile_module
 from torch_tensorrt.dynamo.conversion import convert_module

py/torch_tensorrt/dynamo/lowering/__init__.py

@@ -1,8 +1,5 @@
 from ._decompositions import get_decompositions  # noqa: F401
-from ._fusers import *  # noqa: F403
-from ._pre_aot_lowering import (
-    SUBSTITUTION_REGISTRY,  # noqa: F401
-    register_substitution,  # noqa: F401
-)
-from .substitutions import *  # noqa: F401
 from ._fusers import *  # noqa: F401
+from ._pre_aot_lowering import SUBSTITUTION_REGISTRY  # noqa: F401
+from ._pre_aot_lowering import register_substitution  # noqa: F401
+from .substitutions import *  # noqa: F401

py/torch_tensorrt/dynamo/partitioning/__init__.py

@@ -1,3 +1,3 @@
-from .common import get_submod_inputs
 from ._adjacency_partitioner import partition as fast_partition
 from ._global_partitioner import partition as global_partition
+from .common import get_submod_inputs

py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py

@@ -1,38 +1,35 @@
 import logging
-from typing import Dict, List, Optional, Sequence, Tuple
+from typing import Collection, Dict, List, Optional, Tuple
 
 import torch
-
+import torch.fx.passes.operator_support as ops
+from torch.fx.node import Target
 from torch.fx.passes.splitter_base import (
+    FxNetAccFusionsFinder,
+    FxNetAccNodesFinder,
     Subgraph,
     _SplitterBase,
     _SplitterSettingBase,
-    FxNetAccNodesFinder,
-    FxNetAccFusionsFinder,
 )
-import torch.fx.passes.operator_support as ops
-from torch.fx.passes.tools_common import NodeSet, CALLABLE_NODE_OPS
-from torch.fx.node import Target
-
-from torch_tensorrt.dynamo.conversion.converter_registry import ConverterRegistry
-from .common import DEFAULT_SINGLE_NODE_PARTITIONS
-from torch_tensorrt.dynamo._defaults import MIN_BLOCK_SIZE
-
+from torch.fx.passes.tools_common import CALLABLE_NODE_OPS, NodeSet
 from torch_tensorrt.dynamo import DYNAMO_CONVERTERS as CONVERTERS
+from torch_tensorrt.dynamo._defaults import MIN_BLOCK_SIZE
+from torch_tensorrt.dynamo.conversion.converter_registry import ConverterRegistry
 
+from .common import DEFAULT_SINGLE_NODE_PARTITIONS
 
 logger = logging.getLogger(__name__)
 
 
-class OpSupportTester(ops.OperatorSupportBase):
+class OpSupportTester(ops.OperatorSupportBase):  # type: ignore
     """Class to determine whether operators within a module are supported"""
 
-    def __init__(self, torch_executed_ops: Sequence[Target] = set()) -> None:
+    def __init__(self, torch_executed_ops: Collection[Target] = set()) -> None:
         super().__init__()
 
         # Initialize sets of supported/unsupported operators
-        self.supported_operators = {}
-        self.unsupported_operators = {}
+        self.supported_operators: Dict[str, int] = {}
+        self.unsupported_operators: Dict[str, int] = {}
         self.torch_executed_ops = torch_executed_ops
 
     def is_node_supported(
@@ -58,7 +55,7 @@ def is_node_supported(
 
             return False
 
-    def print_support_overview(self, num_trt_blocks: Optional[int] = None):
+    def print_support_overview(self, num_trt_blocks: Optional[int] = None) -> None:
         if num_trt_blocks is not None:
             logger.debug(
                 f"\nNumber of TensorRT-Accelerated Engines Generated: {num_trt_blocks}"
@@ -81,7 +78,7 @@ def print_support_overview(self, num_trt_blocks: Optional[int] = None):
             logger.debug("\nAll Nodes Supported\n")
 
 
-class TRTPartitioner(_SplitterBase):
+class TRTPartitioner(_SplitterBase):  # type: ignore
     """Partitioner to split an FX graph into subgraphs based on operator support
 
     Adapted from, and modified for the Torch-TensorRT Dynamo case:
@@ -102,7 +99,7 @@ def __init__(
         module: torch.fx.GraphModule,
         operator_support: ops.OperatorSupportBase,
         allowed_single_node_partition_ops: Optional[
-            Sequence[str]
+            Collection[str]
         ] = DEFAULT_SINGLE_NODE_PARTITIONS,
         min_block_size: int = MIN_BLOCK_SIZE,
     ):
@@ -141,7 +138,7 @@ def __init__(
         self.non_acc_submodule_name = "_run_on_gpu_"
         self._node_submodule_map: Dict[str, str] = {}
 
-        self.num_trt_accelerated_subgraphs = None
+        self.num_trt_accelerated_subgraphs: Optional[int] = None
         self.allowed_single_node_partition_ops = allowed_single_node_partition_ops
 
     def remove_small_acc_subgraphs(self, subgraphs: List[Subgraph]) -> List[Subgraph]:
@@ -152,10 +149,13 @@ def remove_small_acc_subgraphs(self, subgraphs: List[Subgraph]) -> List[Subgraph
         result: List[Subgraph] = []
         for subgraph in subgraphs:
             if subgraph.is_acc:
-                if len(subgraph.nodes) >= self.settings.min_acc_module_size or any(
-                    ConverterRegistry.qualified_name_or_str(node.target)
-                    in self.allowed_single_node_partition_ops
-                    for node in subgraph.nodes
+                if len(subgraph.nodes) >= self.settings.min_acc_module_size or (
+                    self.allowed_single_node_partition_ops is not None
+                    and any(
+                        ConverterRegistry.qualified_name_or_str(node.target)
+                        in self.allowed_single_node_partition_ops
+                        for node in subgraph.nodes
+                    )
                 ):
                     result.append(subgraph)
                 else:
@@ -214,7 +214,7 @@ def partition(
     gm: torch.fx.GraphModule,
     verbose: bool = True,
     min_block_size: int = MIN_BLOCK_SIZE,
-    torch_executed_ops: Sequence[Target] = set(),
+    torch_executed_ops: Collection[Target] = set(),
 ) -> torch.fx.GraphModule:
     """Partition an FX GraphModule with aten ops into TRT engines
     Partitioning is based on converter operator support
@@ -223,7 +223,7 @@ def partition(
         gm: FX GraphModule to partition
         verbose: Bool representing whether to print operator support
         min_block_size: Minimum number of operators per TRT-Engine Block
-        torch_executed_ops: Sequence of operations to run in Torch, regardless of converter coverage
+        torch_executed_ops: Collection of operations to run in Torch, regardless of converter coverage
     Returns:
         torch.fx.GraphModule
     """

py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py

@@ -1,20 +1,17 @@
 import logging
-from typing import Dict, List, Optional, Sequence
+from typing import Collection, Dict, List, Mapping, Optional, Sequence, Set
 
 import torch
-
-from torch_tensorrt.dynamo._defaults import MIN_BLOCK_SIZE
-from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.graph_module import GraphModule
-from .common import DEFAULT_SINGLE_NODE_PARTITIONS
 from torch.fx.node import _get_qualified_name
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.passes.operator_support import OperatorSupport, SupportDict
 from torch_tensorrt.dynamo._defaults import MIN_BLOCK_SIZE
 from torch_tensorrt.dynamo.conversion.converter_registry import (
     DYNAMO_CONVERTERS as CONVERTERS,
 )
-from torch_tensorrt.dynamo.lowering._pre_aot_lowering import SUBSTITUTION_REGISTRY
+
+from .common import DEFAULT_SINGLE_NODE_PARTITIONS
 
 logger = logging.getLogger(__name__)
 
@@ -40,7 +37,7 @@ def __init__(
         *,
         non_compute_ops: Optional[Sequence[str]] = None,
         allowed_single_node_partition_ops: Optional[
-            Sequence[str]
+            Collection[str]
         ] = DEFAULT_SINGLE_NODE_PARTITIONS,
         min_block_size: int = MIN_BLOCK_SIZE,
     ) -> None:

py/torch_tensorrt/dynamo/partitioning/common.py

@@ -1,13 +1,14 @@
-import torch
 import logging
-from typing import Sequence, Set
-from torch_tensorrt.dynamo.lowering import SUBSTITUTION_REGISTRY
+from typing import Any, Optional, Sequence, Set
+
+import torch
 from torch.fx.node import _get_qualified_name
+from torch_tensorrt.dynamo.lowering import SUBSTITUTION_REGISTRY
 
-DEFAULT_SINGLE_NODE_PARTITIONS: Set[str] = set(
+DEFAULT_SINGLE_NODE_PARTITIONS: Set[str] = {
     _get_qualified_name(to_replace.new_operator)
     for to_replace in SUBSTITUTION_REGISTRY.values()
-)
+}
 
 
 logger = logging.getLogger(__name__)
@@ -17,7 +18,7 @@ def get_submod_inputs(
     mod: torch.fx.GraphModule,
     submod: torch.fx.GraphModule,
     inputs: Sequence[torch.Tensor],
-) -> Sequence[torch.Tensor]:
+) -> Optional[Sequence[torch.Tensor]]:
     """Helper function to get inputs to a Torch submodule
 
     Args:
@@ -29,9 +30,10 @@ def get_submod_inputs(
     """
     acc_inputs = None
 
-    def get_input(self, inputs):
+    def get_input(self: Any, inputs: Sequence[torch.Tensor]) -> None:
         nonlocal acc_inputs
         acc_inputs = inputs
+        return
 
     handle = submod.register_forward_pre_hook(get_input)
     mod(*inputs)