remove allow-untyped-defs from torch/_prims/executor.py (pytorch#144233)

Pull Request resolved: pytorch#144233
Approved by: https://github.com/Skylion007

Author: bobrenjc93

torch/_prims/executor.py

@@ -1,17 +1,22 @@
-# mypy: allow-untyped-defs
-from typing import Callable, Optional
+from typing import Any, Callable, Optional, TypeVar
+from typing_extensions import ParamSpec, TypeVarTuple, Unpack
 
 from torch._prims.context import TorchRefsMode
 from torch.fx import GraphModule
 from torch.fx.experimental.proxy_tensor import make_fx, wrapper_and_args_for_make_fx
 
 
+T = TypeVar("T")
+P = ParamSpec("P")
+Ts = TypeVarTuple("Ts")
+
+
 def execute(
     gm: GraphModule,
-    *args,
+    *args: Unpack[Ts],
     executor: str = "aten",
     executor_parameters: Optional[dict] = None,
-):
+) -> Any:
     """
     Prototype ATen executor.
 
@@ -25,7 +30,7 @@ def execute(
     raise ValueError(msg)
 
 
-def make_traced(fn: Callable):
+def make_traced(fn: Callable[P, T]) -> Callable[P, T]:
     """
     Returns a function that, when called, will
     trace its torch operations to prims and then
@@ -49,12 +54,14 @@ def foo(a, b):
     result = traced_foo(a, b, executor='aten')
     """
 
-    def _traced(*args, executor="aten", **kwargs):
+    def _traced(*args: P.args, **kwargs: P.kwargs) -> T:
+        executor = str(kwargs.pop("executor", "aten"))
+
         # TODO: caching
         wrapped, all_args = wrapper_and_args_for_make_fx(fn, args, kwargs)
 
         with TorchRefsMode():
             gm = make_fx(wrapped)(all_args)
         return execute(gm, all_args, executor=executor)
 
-    return _traced
+    return _traced  # type: ignore[return-value]