[BE] typing for decorators - distributed/_tensor/ops/utils (pytorch#142139)

Test Plan: unit tests

Differential Revision: D62302679

Pull Request resolved: pytorch#142139
Approved by: https://github.com/Skylion007, https://github.com/kwen2501

Author: aorenste

torch/distributed/tensor/_ops/_conv_ops.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor
 from typing import List

torch/distributed/tensor/_ops/_embedding_ops.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor

torch/distributed/tensor/_ops/_experimental_ops.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor
 

torch/distributed/tensor/_ops/_math_ops.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 import math

torch/distributed/tensor/_ops/_matrix_ops.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor
 

torch/distributed/tensor/_ops/_random_ops.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # Copyright (c) Meta Platforms, Inc. and affiliates
 import torch
 from torch.distributed.device_mesh import DeviceMesh

torch/distributed/tensor/_ops/_tensor_ops.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 from typing import cast, List, Optional, Sequence, Sized, Tuple
@@ -593,7 +592,7 @@ def prop_index_select(op_schema: OpSchema) -> OutputSharding:
             args_schema=(
                 schema_suggestion.args_schema[0],
                 dim,
-                schema_suggestion.args_schema[1][dim],
+                schema_suggestion.args_schema[1][dim],  # type: ignore[index]
             ),
             kwargs_schema=op_schema.kwargs_schema,
         )

torch/distributed/tensor/_ops/_view_ops.py

@@ -1,4 +1,3 @@
-# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 from dataclasses import dataclass

torch/distributed/tensor/_ops/utils.py

@@ -3,7 +3,18 @@
 import functools
 import itertools
 import operator
-from typing import cast, Iterable, List, Optional, Sequence, Tuple, Union
+from typing import (
+    Callable,
+    cast,
+    Iterable,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    TypeVar,
+    Union,
+)
+from typing_extensions import ParamSpec
 
 import torch
 from torch.distributed.tensor._api import DTensor
@@ -25,14 +36,21 @@
 )
 
 
+_T = TypeVar("_T")
+_P = ParamSpec("_P")
+
+
 # convenient wrapper to register sharding propagation rules
 # pyre-fixme[3]: Return type must be annotated.
 # pyre-fixme[2]: Parameter must be annotated.
-def register_prop_rule(op, schema_info=None):
+def register_prop_rule(
+    op: Union[torch._ops.OpOverload, List[torch._ops.OpOverload]],
+    schema_info: Optional[RuntimeSchemaInfo] = None,
+) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]:
     # pyre-fixme[53]: Captured variable `func` is not annotated.
     # pyre-fixme[3]: Return type must be annotated.
     # pyre-fixme[2]: Parameter must be annotated.
-    def wrapper(impl):
+    def wrapper(impl: Callable[_P, _T]) -> Callable[_P, _T]:
         overloads = op if isinstance(op, list) else [op]
         for overload in overloads:
             DTensor._op_dispatcher.sharding_propagator.register_sharding_prop_rule(
@@ -43,7 +61,9 @@ def wrapper(impl):
     return wrapper
 
 
-def register_op_strategy(op, schema_info=None):
+def register_op_strategy(
+    op, schema_info=None
+) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]:
     # pyre-fixme[53]: Captured variable `func` is not annotated.
     # pyre-fixme[3]: Return type must be annotated.
     # pyre-fixme[2]: Parameter must be annotated.