chore(//py/torch_tensorrt/dynamo/lowering): mypy compilance

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

py/torch_tensorrt/dynamo/_settings.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Optional, Sequence
+from typing import Optional, Sequence, Set
 import torch
 from torch_tensorrt.dynamo._defaults import (
     PRECISION,
@@ -20,7 +20,7 @@ class CompilationSettings:
     debug: bool = DEBUG
     workspace_size: int = WORKSPACE_SIZE
     min_block_size: int = MIN_BLOCK_SIZE
-    torch_executed_ops: Sequence[str] = field(default_factory=set)
+    torch_executed_ops: Set[str] = field(default_factory=set)
     pass_through_build_failures: bool = PASS_THROUGH_BUILD_FAILURES
     max_aux_streams: Optional[int] = MAX_AUX_STREAMS
     version_compatible: bool = VERSION_COMPATIBLE

py/torch_tensorrt/dynamo/lowering/_fusers.py

@@ -2,16 +2,24 @@
 from torch_tensorrt.fx.tracer.acc_tracer import acc_ops
 
 
-def check_permute(node: torch.fx.Node):
+def check_permute(node: torch.fx.Node) -> bool:
     ranks = len(node.meta["tensor_meta"].shape)
     permutation = list(i % ranks for i in node.kwargs["permutation"])  # type: ignore[union-attr]
     allowed_permutation = list(i for i in range(ranks))
     allowed_permutation[-1] = ranks - 2
     allowed_permutation[-2] = ranks - 1
     return permutation == allowed_permutation
 
+def trt_transposed_matmul(
+    lhs: torch.Tensor, rhs: torch.Tensor, lhs_transposed: bool, rhs_transposed: bool
+) -> torch.Tensor:
+    if lhs_transposed:
+        lhs = lhs.transpose(-1, -2)
+    if rhs_transposed:
+        rhs = rhs.transpose(-1, -2)
+    return torch.matmul(lhs, rhs)
 
-def fuse_permute_matmul(gm: torch.fx.GraphModule):
+def fuse_permute_matmul(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     """
     Fuse pattern like permute + matmul if permute is transposing the last two dimension.
     """
@@ -45,11 +53,11 @@ def fuse_permute_matmul(gm: torch.fx.GraphModule):
 
 def trt_transposed_linear(
     input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor
-):
+) -> torch.Tensor:
     return torch.matmul(input.transpose(-1, -2), weight.t()) + bias
 
 
-def fuse_permute_linear(gm: torch.fx.GraphModule):
+def fuse_permute_linear(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     """
     Fuse pattern like permute + linear if permute is transposing the last two dimension.
     """

py/torch_tensorrt/dynamo/lowering/_partition.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Dict, List, Optional, Sequence, Set
+from typing import Dict, List, Optional, Sequence, Set, Mapping, Any
 
 import torch
 
@@ -8,14 +8,14 @@
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.graph_module import GraphModule
 from torch.fx.node import _get_qualified_name
-from torch.fx.passes.operator_support import OperatorSupport
+from torch.fx.passes.operator_support import OperatorSupport, SupportDict
 
 from torch_tensorrt.fx.converter_registry import CONVERTERS
 
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_SINGLE_NODE_PARTITIONS: Set[str] = set(
+DEFAULT_SINGLE_NODE_PARTITIONS: List[str]  = list(
     _get_qualified_name(to_replace.new_operator)
     for to_replace in SUBSTITUTION_REGISTRY.values()
 )
@@ -41,10 +41,8 @@ def __init__(
         operator_support: OperatorSupport,
         *,
         non_compute_ops: Optional[Sequence[str]] = None,
-        allowed_single_node_partition_ops: Optional[
-            Sequence[str]
-        ] = DEFAULT_SINGLE_NODE_PARTITIONS,
-        min_block_size=MIN_BLOCK_SIZE,
+        allowed_single_node_partition_ops: Optional[Sequence[str]] = DEFAULT_SINGLE_NODE_PARTITIONS,
+        min_block_size: int = MIN_BLOCK_SIZE,
     ) -> None:
         super().__init__(
             graph_module,
@@ -74,14 +72,14 @@ def propose_partitions(self) -> List[Partition]:
                 # Partitions are exempted from min_block_size if they contain an allowed single-node op
                 if (
                     node.op == "call_function"
-                    and _get_qualified_name(node.target)
+                    and _get_qualified_name(node.target) # type: ignore[arg-type]
                     in self.allowed_single_node_partition_ops
                 ):
                     exempted_partition = True
                     break
                 elif (
                     node.op == "call_function"
-                    and _get_qualified_name(node.target) not in non_compute_ops
+                    and _get_qualified_name(node.target) not in non_compute_ops # type: ignore[arg-type]
                 ):
                     compute_node_count += 1
 
@@ -106,16 +104,16 @@ def partition_and_fuse(self) -> GraphModule:
 class TorchTensorRTOperatorSupport(OperatorSupport):
     """Class to determine whether operators within a module are supported"""
 
-    def __init__(self, support_dict=None, torch_executed_ops=set()):
+    def __init__(self, support_dict: Optional[SupportDict] = None, torch_executed_ops: Set[str] = set()):
         super().__init__(support_dict)
 
         # Initialize sets of supported/unsupported operators
-        self.supported_operators = set()
-        self.unsupported_operators = set()
-        self.torch_executed_ops = torch_executed_ops
+        self.supported_operators: Set[str] = set()
+        self.unsupported_operators: Set[str]  = set()
+        self.torch_executed_ops: Set[str] = torch_executed_ops
 
     def is_node_supported(
-        self, submodules: Dict[str, torch.nn.Module], node: torch.fx.Node
+        self, submodules: Mapping[str, torch.nn.Module], node: torch.fx.Node
     ) -> bool:
         node_name = (
             _get_qualified_name(node.target)
@@ -138,7 +136,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}"
@@ -164,7 +162,7 @@ def partition(
     gm: torch.fx.GraphModule,
     verbose: bool = True,
     min_block_size: int = MIN_BLOCK_SIZE,
-    torch_executed_ops: Sequence[str] = set(),
+    torch_executed_ops: Set[str] = set(),
 ) -> torch.fx.GraphModule:
     """Partition an FX GraphModule with aten ops into TRT engines
     Partitioning is based on converter operator support
@@ -195,7 +193,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:
@@ -205,9 +203,9 @@ def get_submod_inputs(
     Returns:
         Sequence of Tensors representing inputs to child module
     """
-    acc_inputs = None
+    acc_inputs: Optional[Sequence[torch.Tensor]] = None
 
-    def get_input(self, inputs):
+    def get_input(_: torch.fx.GraphModule, inputs: Sequence[torch.Tensor]) -> None:
         nonlocal acc_inputs
         acc_inputs = inputs
 

py/torch_tensorrt/dynamo/lowering/substitutions/einsum.py

@@ -1,4 +1,4 @@
-from typing import Dict, Tuple
+from typing import Dict, Tuple, Any, List, Optional, Sequence
 import torch
 from torch._custom_op.impl import custom_op
 from torch.fx.node import Argument, Target
@@ -14,18 +14,18 @@
     qualname="tensorrt::einsum",
     manual_schema="(str equation, Tensor[] tensors) -> Tensor",
 )
-def einsum(equation, tensors):
+def einsum(equation: str, tensors: List[torch.Tensor]) -> torch.Tensor: # type: ignore[empty-body]
     # Defines operator schema, name, namespace, and function header
     ...
 
 
-@einsum.impl("cpu")
-@einsum.impl("cuda")
-@einsum.impl_abstract()
+@einsum.impl("cpu") # type: ignore[misc]
+@einsum.impl("cuda") # type: ignore[misc]
+@einsum.impl_abstract() # type: ignore[misc]
 def einsum_generic(
-    *args,
-    **kwargs,
-):
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
     # Defines a converter implementation for AOT Autograd to use for shape analysis/propagation
     return torch.einsum(
         *args,
@@ -42,6 +42,7 @@ def aten_ops_einsum(
     name: str,
 ) -> TRTTensor:
     # Defines converter replacing the default operator for this function
+    assert isinstance(args[1], Sequence)
     for input_trt in args[1]:
         if not isinstance(input_trt, TRTTensor):
             raise RuntimeError(f"Einsum received non-TRTTensor input: {input_trt}")
@@ -56,7 +57,7 @@ def aten_ops_einsum(
 def einsum_insertion_fn(
     gm: torch.fx.GraphModule,
     node: torch.fx.Node,
-    _unused: None = None,
+    _: Optional[torch.nn.Module] = None,
 ) -> torch.fx.Node:
     equation = node.args[0]
 
@@ -71,7 +72,7 @@ def einsum_insertion_fn(
     ), f"TRT Einsum currently only supports 1 or 2 Tensors, got {len(inputs)} Tensors"
 
     # Ensure the input is formatted as an equation and
-    new_node = gm.graph.call_function(
+    new_node: torch.fx.Node = gm.graph.call_function(
         torch.ops.tensorrt.einsum,
         args=(equation, inputs),
         kwargs=node.kwargs,

py/torch_tensorrt/dynamo/lowering/substitutions/maxpool1d.py

@@ -1,4 +1,4 @@
-from typing import Dict, Tuple
+from typing import Dict, Tuple, Any, Optional
 import torch
 from torch._custom_op.impl import custom_op
 from torch.fx.node import Argument, Target
@@ -26,7 +26,7 @@
     qualname="tensorrt::maxpool1d",
     manual_schema="(Tensor x, int[1] kernel_size, int[1] stride, int[1] padding, int[1] dilation, bool ceil_mode) -> Tensor",
 )
-def maxpool1d(x, kernel_size, stride, padding, dilation, ceil_mode):
+def maxpool1d(x: torch.Tensor, kernel_size: Tuple[int], stride: Tuple[int], padding: Tuple[int], dilation: Tuple[int], ceil_mode: bool) -> torch.Tensor: # type: ignore[empty-body]
     # Defines operator schema, name, namespace, and function header
     ...
 
@@ -38,13 +38,13 @@ def maxpool1d(x, kernel_size, stride, padding, dilation, ceil_mode):
 # is desirable. If the operator to replace is a custom module you've written, then add its Torch
 # implementation here. Note that the function header to the generic function can have specific arguments
 # as in the above placeholder
-@maxpool1d.impl("cpu")
-@maxpool1d.impl("cuda")
-@maxpool1d.impl_abstract()
+@maxpool1d.impl("cpu") # type: ignore[misc]
+@maxpool1d.impl("cuda") # type: ignore[misc]
+@maxpool1d.impl_abstract() # type: ignore[misc]
 def maxpool1d_generic(
-    *args,
-    **kwargs,
-):
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
     # Defines an implementation for AOT Autograd to use for shape analysis/propagation
     return torch.nn.functional.max_pool1d(
         *args,
@@ -75,10 +75,11 @@ def maxpool1d_generic(
 def maxpool1d_insertion_fn(
     gm: torch.fx.GraphModule,
     node: torch.fx.Node,
-    submodule: torch.nn.Module,
+    submodule: Optional[torch.nn.Module],
 ) -> torch.fx.Node:
     # Defines insertion function for new node
-    new_node = gm.graph.call_function(
+    assert submodule is not None
+    new_node: torch.fx.Node = gm.graph.call_function(
         torch.ops.tensorrt.maxpool1d,
         args=node.args,
         kwargs={

pyproject.toml

@@ -83,6 +83,11 @@ exclude = [
     "docsrc/",
     "tests/",
 ]
+python_version = "3.10"
+
+[[tool.mypy.overrides]]
+module = "torch_tensorrt.dynamo.lowering._decompositions"
+disallow_untyped_calls = false
 
 [tool.setuptools]
 package-dir = {"" = "py"}