chore(//py/torch_tensorrt): subsitutions mypy compliant

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

Author: narendasan

py/torch_tensorrt/dynamo/lowering/_decompositions.py

@@ -3,7 +3,7 @@
 from torch._decomp import register_decomposition, core_aten_decompositions, OpOverload
 
 
-DECOMPOSITIONS: Dict[OpOverload, Callable[[Any], Any]] = {**core_aten_decompositions()}
+DECOMPOSITIONS: Dict[OpOverload, Callable[..., Any]] = {**core_aten_decompositions()}
 
 aten = torch.ops.aten
 
@@ -60,5 +60,5 @@ def addmm_replacement(
     )
 
 
-def get_decompositions() ->  Dict[OpOverload, Callable[[Any], Any]]:
+def get_decompositions() ->  Dict[OpOverload, Callable[..., Any]]:
     return DECOMPOSITIONS

py/torch_tensorrt/dynamo/lowering/_pre_aot_lowering.py

@@ -6,6 +6,10 @@
 
 logger = logging.getLogger(__name__)
 
+SubgraphInsertionFnType = Callable[
+    [torch.fx.GraphModule, torch.fx.Node, Optional[torch.nn.Module]], torch.fx.Node
+]
+
 
 @dataclass(frozen=True)
 class Substitution:
@@ -18,22 +22,20 @@ class Substitution:
     # and returning a replacement node, with type 'call_function', or raising an Error if
     # incompatibility is detected
     # Note: subgraph_insertion_fn should NOT delete nodes or recompile the graph
-    subgraph_insertion_fn: Callable[
-        [torch.fx.GraphModule, torch.fx.Node, Optional[torch.nn.Module]], torch.fx.Node
-    ]
+    subgraph_insertion_fn: SubgraphInsertionFnType
 
 
 # Dictionary mapping module to Substitution instance
 SUBSTITUTION_REGISTRY: Dict[
-    Union[Type[torch.nn.Module], Callable], Substitution
+    Union[Type[torch.nn.Module], Callable[..., Any]], Substitution
 ] = dict()
 
 
 def register_substitution(
-    module_or_function_to_replace: Union[Type[torch.nn.Module], Callable],
+    module_or_function_to_replace: Union[Type[torch.nn.Module], Callable[..., Any]],
     new_operator: torch._ops.OpOverload,
     enabled: bool = True,
-) -> Callable[[Any], Any]:
+) -> Callable[[SubgraphInsertionFnType], SubgraphInsertionFnType]:
     """Decorator to register subgraph insertion functions
 
     Args:
@@ -44,22 +46,22 @@ def register_substitution(
         torch.fx.GraphModule
     """
 
-    def enable_substitution(subgraph_insertion_fn):
+    def enable_substitution(subgraph_insertion_fn: SubgraphInsertionFnType) -> SubgraphInsertionFnType:
         """Function for use if substitution is enabled"""
         replacement = Substitution(
             new_operator=new_operator, subgraph_insertion_fn=subgraph_insertion_fn
         )
         SUBSTITUTION_REGISTRY[module_or_function_to_replace] = replacement
         return subgraph_insertion_fn
 
-    def disable_substitution(subgraph_insertion_fn):
+    def disable_substitution(subgraph_insertion_fn: SubgraphInsertionFnType) -> SubgraphInsertionFnType:
         """Function for use if substitution is disabled"""
         return subgraph_insertion_fn
 
     return enable_substitution if enabled else disable_substitution
 
 
-def pre_aot_substitutions(gm: torch.fx.GraphModule):
+def pre_aot_substitutions(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     """Perform graph substitutions prior to AOT tracing
 
     Args:
@@ -92,6 +94,7 @@ def pre_aot_substitutions(gm: torch.fx.GraphModule):
         # If submodule/function is a member of the substitution registry, replace it
         if exists_in_registry:
             try:
+                assert to_replace is not None
                 replacement = SUBSTITUTION_REGISTRY[to_replace]
                 op, insertion_fn = (
                     replacement.new_operator,