chore(//py/torch_tensorrt/dynamo): Tracer mypy compliance

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

Author: narendasan

py/torch_tensorrt/dynamo/aten_tracer.py

@@ -6,8 +6,10 @@
 
 import torch
 import torch._dynamo as torchdynamo
+from torch import _guards
+from torch.fx.passes.infra.pass_base import PassResult
 
-from torch_tensorrt.fx.utils import req_torch_version
+from torch_tensorrt.dynamo.utils import req_torch_version
 from torch_tensorrt.fx.passes.lower_basic_pass_aten import (
     compose_bmm,
     compose_chunk,
@@ -97,7 +99,7 @@ def dynamo_trace(
     aten_graph: bool,
     tracing_mode: str = "real",
     dynamo_config: Optional[DynamoConfig] = None,
-) -> Tuple[torch.fx.GraphModule, Set]:
+) -> Tuple[torch.fx.GraphModule, Set[_guards.Guard]]:
     """
     TODO: Once we fully migrate to torchdynamo frontend, we will remove
     this config option alltogether.  For now, it helps with quick
@@ -126,7 +128,7 @@ def dynamo_trace(
 
 
 @req_torch_version("2.dev")
-def trace(model, inputs, **kwargs):
+def trace(model: Union[torch.nn.Module, torch.fx.GraphModule], inputs: Tuple[Any, ...], **kwargs: Any) -> torch.fx.GraphModule:
     """
     Optimized trace with necessary passes which re-compose some ops or replace some ops
     These passes should be general and functional purpose
@@ -147,11 +149,11 @@ def trace(model, inputs, **kwargs):
     fx_module, __package__ = dynamo_trace(model, inputs, True, "symbolic")
     print(fx_module.graph)
     for passes in passes_list:
-        pr: PassResult = passes(fx_module)
+        pr: PassResult = passes(fx_module) #type: ignore[assignment] #The type hints in fx are wrong
         fx_module = pr.graph_module
 
     fx_module(*inputs)
 
     fx_module = run_const_fold(fx_module)
     print(fx_module.graph)
-    return fx_module
+    return fx_module #type: ignore[no-any-return]

py/torch_tensorrt/dynamo/backend/backends.py

@@ -139,6 +139,8 @@ def _compile_module(
             partitioned_module, submodule, sample_inputs
         )
 
+        assert submodule_inputs is not None
+
         # Create TRT Module from submodule
         trt_mod = convert_module(
             submodule,

py/torch_tensorrt/dynamo/utils.py

@@ -4,7 +4,8 @@
 from torch_tensorrt.dynamo import CompilationSettings
 from typing import Any, Union, Sequence, Dict
 from torch_tensorrt import Input, Device
-from typing import Optional
+from typing import Optional, Callable, Any
+from packaging import version
 
 logger = logging.getLogger(__name__)
 
@@ -160,3 +161,35 @@ def parse_dynamo_kwargs(kwargs: Dict) -> CompilationSettings:
     logger.debug(f"Compiling with Settings:\n{settings}")
 
     return settings
+
+def req_torch_version(min_torch_version: str = "2.dev") -> Callable[..., Any]:
+    """
+    Create a decorator which verifies the Torch version installed
+    against a specified version range
+
+    Args:
+        min_torch_version (str): The minimum required Torch version
+        for the decorated function to work properly
+
+    Returns:
+        A decorator which raises a descriptive error message if
+        an unsupported Torch version is used
+    """
+
+    def nested_decorator(f: Callable[..., Any]) -> Callable[..., Any]:
+        def function_wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
+            # Parse minimum and current Torch versions
+            min_version = version.parse(min_torch_version)
+            current_version = version.parse(torch.__version__)
+
+            if current_version < min_version:
+                raise AssertionError(
+                    f"Expected Torch version {min_torch_version} or greater, "
+                    + f"when calling {f}. Detected version {torch.__version__}"
+                )
+            else:
+                return f(*args, **kwargs)
+
+        return function_wrapper
+
+    return nested_decorator