chore(//py/torch_tensorrt/dynamo/backend): Backend is mypy conforming

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

Author: narendasan

py/torch_tensorrt/dynamo/backend/backends.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Sequence
+from typing import Sequence, Any, Callable
 import torch
 from functools import partial
 import torch._dynamo as td
@@ -24,19 +24,23 @@
 logger = logging.getLogger(__name__)
 
 
-@td.register_backend(name="torch_tensorrt")
+@td.register_backend(name="torch_tensorrt") # type: ignore[misc]
 def torch_tensorrt_backend(
-    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor], **kwargs
-):
+    gm: torch.fx.GraphModule,
+    sample_inputs: Sequence[torch.Tensor],
+    **kwargs: Any
+) -> torch.nn.Module:
     DEFAULT_BACKEND = aot_torch_tensorrt_aten_backend
 
-    return DEFAULT_BACKEND(gm, sample_inputs, **kwargs)
-
+    compiled_mod: torch.nn.Module = DEFAULT_BACKEND(gm, sample_inputs, **kwargs)
+    return compiled_mod
 
-@td.register_backend(name="aot_torch_tensorrt_aten")
+@td.register_backend(name="aot_torch_tensorrt_aten") # type: ignore[misc]
 def aot_torch_tensorrt_aten_backend(
-    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor], **kwargs
-):
+    gm: torch.fx.GraphModule,
+    sample_inputs: Sequence[torch.Tensor],
+    **kwargs: Any
+) -> torch.nn.Module:
     settings = parse_dynamo_kwargs(kwargs)
 
     custom_backend = partial(
@@ -51,7 +55,7 @@ def aot_torch_tensorrt_aten_backend(
     return aot_module_simplified(
         gm,
         sample_inputs,
-        fw_compiler=make_boxed_compiler(custom_backend),
+        fw_compiler=make_boxed_compiler(custom_backend), # type: ignore[no-untyped-call]
         decompositions=get_decompositions(),
     )
 
@@ -60,7 +64,7 @@ def _pretraced_backend(
     gm: torch.fx.GraphModule,
     sample_inputs: Sequence[torch.Tensor],
     settings: CompilationSettings = CompilationSettings(),
-):
+) -> torch.fx.GraphModule | Callable[..., Any]:
     """Helper function to manage translation of traced FX module to TRT engines
 
     Args:

py/torch_tensorrt/dynamo/lowering/_decompositions.py

@@ -1,20 +1,20 @@
+from typing import Callable, Dict, Any
 import torch
-from torch._decomp import register_decomposition, core_aten_decompositions
+from torch._decomp import register_decomposition, core_aten_decompositions, OpOverload
 
 
-DECOMPOSITIONS = {**core_aten_decompositions()}
+DECOMPOSITIONS: Dict[OpOverload, Callable[[Any], Any]] = {**core_aten_decompositions()}
 
 aten = torch.ops.aten
 
-
-def replace_inplace_op(aten_op, outplace_op):
+def replace_inplace_op(aten_op: OpOverload, outplace_op: OpOverload) -> Any:
     """Replace inplace operation with functional equivalent
     Adapted from:
     https://github.com/pytorch/pytorch/blob/3344d79e3f732dadd5c85b99a7aa1a022f187929/torch/_decomp/decompositions.py#L3355-L3361
     """
 
-    @register_decomposition(aten_op, registry=DECOMPOSITIONS)
-    def inplace_op(*args, **kwargs):
+    @register_decomposition(aten_op, registry=DECOMPOSITIONS) # type: ignore[misc]
+    def inplace_op(*args: Any, **kwargs: Any) -> Any:
         out = outplace_op(*args, **kwargs)
         return args[0].copy_(out)
 
@@ -36,29 +36,29 @@ def inplace_op(*args, **kwargs):
 replace_inplace_op(aten.scatter_reduce_, aten.scatter_reduce)
 
 
-@register_decomposition(aten.std, registry=DECOMPOSITIONS)
-def std_replacement(*args, **kwargs) -> torch.Tensor:
+@register_decomposition(aten.std, registry=DECOMPOSITIONS)  # type: ignore[misc]
+def std_replacement(*args: Any, **kwargs: Any) -> torch.Tensor:
     return torch.sqrt(torch.var(*args, **kwargs))
 
 
-@register_decomposition(aten.rsqrt, registry=DECOMPOSITIONS)
-def rsqrt_replacement(*args, **kwargs) -> torch.Tensor:
+@register_decomposition(aten.rsqrt, registry=DECOMPOSITIONS)  # type: ignore[misc]
+def rsqrt_replacement(*args: Any, **kwargs: Any) -> torch.Tensor:
     return torch.reciprocal(torch.sqrt(*args, **kwargs))
 
 
-@register_decomposition(aten.alias, registry=DECOMPOSITIONS)
+@register_decomposition(aten.alias, registry=DECOMPOSITIONS)  # type: ignore[misc]
 def alias_replacement(x: torch.Tensor) -> torch.Tensor:
     return x
 
 
-@register_decomposition(torch.ops.aten.addmm, registry=DECOMPOSITIONS)
+@register_decomposition(torch.ops.aten.addmm, registry=DECOMPOSITIONS)  # type: ignore[misc]
 def addmm_replacement(
-    input_: torch.Tensor, mat1: torch.Tensor, mat2: torch.Tensor, *, beta=1, alpha=1
+    input_: torch.Tensor, mat1: torch.Tensor, mat2: torch.Tensor, *, beta: int = 1, alpha: int = 1
 ) -> torch.Tensor:
     return torch.add(
         torch.mul(input_, beta), torch.mul(torch.matmul(mat1, mat2), alpha)
     )
 
 
-def get_decompositions():
+def get_decompositions() ->  Dict[OpOverload, Callable[[Any], Any]]:
     return DECOMPOSITIONS