PEP585 update - mostly toplevels (pytorch#145178)

See pytorch#145101 for details.

Pull Request resolved: pytorch#145178
Approved by: https://github.com/bobrenjc93

Author: aorenste

torch/_C/_cudnn.pyi

@@ -1,12 +1,12 @@
 from enum import Enum
 
-from torch.types import _bool, Tuple
+from torch.types import _bool
 
 # Defined in torch/csrc/cuda/shared/cudnn.cpp
 is_cuda: _bool
 
-def getRuntimeVersion() -> Tuple[int, int, int]: ...
-def getCompileVersion() -> Tuple[int, int, int]: ...
+def getRuntimeVersion() -> tuple[int, int, int]: ...
+def getCompileVersion() -> tuple[int, int, int]: ...
 def getVersionInt() -> int: ...
 
 class RNNMode(int, Enum):

torch/__init__.py

@@ -24,13 +24,9 @@
 from typing import (
     Any as _Any,
     Callable as _Callable,
-    Dict as _Dict,
     get_origin as _get_origin,
     Optional as _Optional,
     overload as _overload,
-    Set as _Set,
-    Tuple as _Tuple,
-    Type as _Type,
     TYPE_CHECKING,
     TypeVar as _TypeVar,
     Union as _Union,
@@ -337,7 +333,7 @@ def _load_global_deps() -> None:
     except OSError as err:
         # Can only happen for wheel with cuda libs as PYPI deps
         # As PyTorch is not purelib, but nvidia-*-cu12 is
-        cuda_libs: _Dict[str, str] = {
+        cuda_libs: dict[str, str] = {
             "cublas": "libcublas.so.*[0-9]",
             "cudnn": "libcudnn.so.*[0-9]",
             "cuda_nvrtc": "libnvrtc.so.*[0-9]",
@@ -586,7 +582,7 @@ def __hash__(self) -> builtins.int:
             # https://github.com/arogozhnikov/einops/blob/6181e1e95dc58c00a3143c1726da1c6ee0463164/einops/einops.py#L237
             # return hash(builtins.int(self))
 
-    def as_integer_ratio(self) -> _Tuple["SymInt", builtins.int]:
+    def as_integer_ratio(self) -> tuple["SymInt", builtins.int]:
         """Represent this int as an exact integer ratio"""
         return self, 1
 
@@ -698,7 +694,7 @@ def is_integer(self):
         """Return True if the float is an integer."""
         raise TypeError("type stub not overridden")
 
-    def as_integer_ratio(self) -> _Tuple[builtins.int, builtins.int]:
+    def as_integer_ratio(self) -> tuple[builtins.int, builtins.int]:
         """Represent this float as an exact integer ratio"""
         return builtins.float(self).as_integer_ratio()
 
@@ -857,22 +853,22 @@ def sym_max(a, b):
     assert isinstance(a, all_types), type(a)
     assert isinstance(b, all_types), type(b)
     if isinstance(a, float_types) or isinstance(b, float_types):
-        return builtins.float(builtins.max(a, b))
+        return builtins.float(builtins.max(a, b))  # type: ignore[call-overload]
     else:
-        return builtins.max(a, b)
+        return builtins.max(a, b)  # type: ignore[call-overload]
 
 
-def __all_and_float_types() -> _Tuple[_Tuple[_Type, ...], _Tuple[_Type, ...]]:
+def __all_and_float_types() -> tuple[tuple[type, ...], tuple[type, ...]]:
     try:
         import numpy as np
 
-        all_types: _Tuple[_Type, ...] = (
+        all_types: tuple[type, ...] = (
             np.integer,
             np.floating,
             builtins.int,
             builtins.float,
         )
-        float_types: _Tuple[_Type, ...] = (np.floating, builtins.float)
+        float_types: tuple[type, ...] = (np.floating, builtins.float)
     except ModuleNotFoundError:
         all_types = (builtins.int, builtins.float)
         float_types = (builtins.float,)
@@ -894,9 +890,9 @@ def sym_min(a, b):
     assert isinstance(a, all_types), type(a)
     assert isinstance(b, all_types), type(b)
     if isinstance(a, float_types) or isinstance(b, float_types):
-        return builtins.float(builtins.min(a, b))
+        return builtins.float(builtins.min(a, b))  # type: ignore[call-overload]
     else:
-        return builtins.min(a, b)
+        return builtins.min(a, b)  # type: ignore[call-overload]
 
 
 def sym_sum(args):
@@ -1204,7 +1200,7 @@ def set_default_device(
     _GLOBAL_DEVICE_CONTEXT.device_context = device_context
 
 
-def set_default_tensor_type(t: _Union[_Type["torch.Tensor"], str], /) -> None:
+def set_default_tensor_type(t: _Union[type["torch.Tensor"], str], /) -> None:
     r"""
     .. warning::
 
@@ -2007,7 +2003,7 @@ def _dtype(self):
         return torch.quint2x4
 
 
-_storage_classes: _Set[_Type[_Union[TypedStorage, UntypedStorage]]] = {
+_storage_classes: set[type[_Union[TypedStorage, UntypedStorage]]] = {
     UntypedStorage,
     DoubleStorage,
     FloatStorage,
@@ -2030,7 +2026,7 @@ def _dtype(self):
 }
 
 # The _tensor_classes set is initialized by the call to initialize_python_bindings.
-_tensor_classes: _Set[_Type["torch.Tensor"]] = set()
+_tensor_classes: set[type["torch.Tensor"]] = set()
 
 # If you edit these imports, please update torch/__init__.py.in as well
 from torch import amp as amp, random as random, serialization as serialization
@@ -2282,7 +2278,7 @@ class _TorchCompileInductorWrapper:
     def __init__(self, mode, options, dynamic):
         from torch._inductor.compiler_bisector import CompilerBisector
 
-        self.config: _Dict[str, _Any] = {}
+        self.config: dict[str, _Any] = {}
         self.dynamic = dynamic
         self.apply_mode(mode)
         self.apply_options(options)
@@ -2309,13 +2305,13 @@ def apply_mode(self, mode: _Optional[str]):
 
             self.apply_options(list_mode_options(mode, self.dynamic))
 
-    def apply_options(self, options: _Optional[_Dict[str, _Any]]):
+    def apply_options(self, options: _Optional[dict[str, _Any]]):
         if not options:
             return
 
         from torch._inductor import config
 
-        current_config: _Dict[str, _Any] = config.get_config_copy()
+        current_config: dict[str, _Any] = config.get_config_copy()
 
         for key, val in options.items():
             attr_name = key.replace("-", "_")
@@ -2403,7 +2399,7 @@ def compile(
     dynamic: _Optional[builtins.bool] = None,
     backend: _Union[str, _Callable] = "inductor",
     mode: _Union[str, None] = None,
-    options: _Optional[_Dict[str, _Union[str, builtins.int, builtins.bool]]] = None,
+    options: _Optional[dict[str, _Union[str, builtins.int, builtins.bool]]] = None,
     disable: builtins.bool = False,
 ) -> _Callable[_InputT, _RetT]: ...
 
@@ -2416,7 +2412,7 @@ def compile(
     dynamic: _Optional[builtins.bool] = None,
     backend: _Union[str, _Callable] = "inductor",
     mode: _Union[str, None] = None,
-    options: _Optional[_Dict[str, _Union[str, builtins.int, builtins.bool]]] = None,
+    options: _Optional[dict[str, _Union[str, builtins.int, builtins.bool]]] = None,
     disable: builtins.bool = False,
 ) -> _Callable[[_Callable[_InputT, _RetT]], _Callable[_InputT, _RetT]]: ...
 
@@ -2428,7 +2424,7 @@ def compile(
     dynamic: _Optional[builtins.bool] = None,
     backend: _Union[str, _Callable] = "inductor",
     mode: _Union[str, None] = None,
-    options: _Optional[_Dict[str, _Union[str, builtins.int, builtins.bool]]] = None,
+    options: _Optional[dict[str, _Union[str, builtins.int, builtins.bool]]] = None,
     disable: builtins.bool = False,
 ) -> _Union[
     _Callable[[_Callable[_InputT, _RetT]], _Callable[_InputT, _RetT]],
@@ -2624,7 +2620,7 @@ def _register_device_module(device_type, module):
 
     class _TritonLibrary:
         lib = torch.library.Library("triton", "DEF")
-        ops_table: _Dict[_Tuple[str, str], _Callable] = {}
+        ops_table: dict[tuple[str, str], _Callable] = {}
 
         @classmethod
         def registerOp(cls, op_key, full_schema, op_impl, dispatch_key):

torch/_custom_op/impl.py

@@ -108,7 +108,7 @@ def inner(func):
 # An example usage is FakeTensor: FakeTensor checks if a specific operator
 # has an implementation registered via the CustomOp API.
 # Indexed by qualname (e.g. aten::foo)
-global_registry: typing.Dict[str, "CustomOp"] = {}
+global_registry: dict[str, "CustomOp"] = {}
 
 
 class CustomOp:
@@ -136,7 +136,7 @@ def __init__(self, lib, cpp_ns, schema, operator_name, ophandle, *, _private_acc
         self.__name__ = None  # mypy requires this
         # NB: Some of these impls are registered as kernels to DispatchKeys.
         # Modifying the _impls dict directly won't do anything in that case.
-        self._impls: typing.Dict[str, typing.Optional[FuncAndLocation]] = {}
+        self._impls: dict[str, typing.Optional[FuncAndLocation]] = {}
         # See NOTE [CustomOp autograd kernel indirection]
         self._registered_autograd_kernel_indirection = False
 
@@ -476,7 +476,7 @@ def validate_schema(schema: FunctionSchema) -> None:
         )
 
 
-def parse_qualname(qualname: str) -> typing.Tuple[str, str]:
+def parse_qualname(qualname: str) -> tuple[str, str]:
     names = qualname.split("::", 1)
     if len(names) != 2:
         raise ValueError(f"Expected there to be a namespace in {qualname}, i.e. The "

torch/_dispatch/python.py

@@ -1,8 +1,8 @@
 # mypy: allow-untyped-defs
 import itertools
 import unittest.mock
+from collections.abc import Iterator
 from contextlib import contextmanager
-from typing import Iterator
 
 import torch
 import torch._C

torch/_guards.py

@@ -17,13 +17,9 @@
 from typing import (
     Any,
     Callable,
-    Dict,
     Generic,
-    List,
     NamedTuple,
     Optional,
-    Set,
-    Tuple,
     TYPE_CHECKING,
     TypeVar,
     Union,
@@ -260,8 +256,8 @@ class Guard:
     create_fn: Callable[[GuardBuilderBase, Guard], None]
 
     # Export only. These values are written to at time of guard check_fn creation.
-    guard_types: Optional[List[str]] = None
-    code_list: Optional[List[str]] = None
+    guard_types: Optional[list[str]] = None
+    code_list: Optional[list[str]] = None
     obj_weakref: Optional[object] = None
     guarded_class_weakref: Optional[type] = None
 
@@ -448,8 +444,8 @@ def __post_init__(self):
 
 @dataclasses.dataclass
 class StorageOverlap(GuardEnvExpr):
-    overlapping_sources: List[Source]
-    non_overlapping_sources: List[Source]
+    overlapping_sources: list[Source]
+    non_overlapping_sources: list[Source]
 
 
 """
@@ -478,7 +474,7 @@ class GuardsCheckpointState:
     The GuardCheckpointState - it is the T of Checkpointable[T] for GuardsContext
     """
 
-    dynamo_guards: Set[Guard] = set()
+    dynamo_guards: set[Guard] = set()
 
     def __init__(self, dynamo_guards):
         self.dynamo_guards = dynamo_guards
@@ -500,7 +496,7 @@ def __eq__(self, other):
 
 
 class ModuleContextCheckpointState:
-    nn_modules: Dict[str, torch.nn.Module] = {}
+    nn_modules: dict[str, torch.nn.Module] = {}
 
     def __init__(self, nn_modules):
         self.nn_modules = nn_modules
@@ -523,7 +519,7 @@ def __eq__(self, other):
 
 class ModuleContext(Checkpointable[ModuleContextCheckpointState]):
     def __init__(self) -> None:
-        self.nn_modules: Dict[str, Any] = {}
+        self.nn_modules: dict[str, Any] = {}
 
     def copy_graphstate(self):
         return ModuleContextCheckpointState(dict(self.nn_modules))
@@ -534,7 +530,7 @@ def restore_graphstate(self, state):
 
 
 class GlobalContextCheckpointState:
-    global_state: Dict[str, Tuple[Callable, ...]] = {}
+    global_state: dict[str, tuple[Callable, ...]] = {}
 
     def __init__(self, global_states):
         self.global_state = global_states
@@ -572,7 +568,7 @@ class GlobalContext(Checkpointable[GlobalContextCheckpointState]):
     }
 
     def __init__(self) -> None:
-        self.global_state: Dict[str, Tuple[Callable, ...]] = {}
+        self.global_state: dict[str, tuple[Callable, ...]] = {}
 
     def copy_graphstate(self):
         return GlobalContextCheckpointState(dict(self.global_state))
@@ -628,7 +624,7 @@ def add(self, guard: Guard, *, collect_debug_stack=True, skip=0):
                 guard.user_stack = TracingContext.extract_stack()
         self.inner.add(guard)
 
-    def update(self, *others: Set[Guard]):
+    def update(self, *others: set[Guard]):
         for o in others:
             for g in o:
                 self.add(g, skip=1)
@@ -641,7 +637,7 @@ def remove_guards_with_source(self, source):
 class GuardsContext(Checkpointable[GuardsCheckpointState]):
     def __init__(self) -> None:
         self.dynamo_guards: GuardsSet = GuardsSet()
-        self.aotautograd_guards: List[GuardEnvExpr] = []
+        self.aotautograd_guards: list[GuardEnvExpr] = []
 
     def copy_graphstate(self):
         return GuardsCheckpointState(set(self.dynamo_guards.inner))
@@ -674,9 +670,9 @@ def get_proxy_dispatch_entry(self, identifier: str): ...
 
 class InvokeSubgraphCache(HopSubgraphCache):
     def __init__(self) -> None:
-        self.autograd_cache: Dict[str, Callable] = {}
-        self.proxy_dispatch_cache: Dict[str, Callable] = {}
-        self.dynamo_identifiers: Dict[str, str] = {}
+        self.autograd_cache: dict[str, Callable] = {}
+        self.proxy_dispatch_cache: dict[str, Callable] = {}
+        self.dynamo_identifiers: dict[str, str] = {}
 
     def add_dynamo_identifier(self, cache_key: str, identifier: str):
         self.dynamo_identifiers[cache_key] = identifier
@@ -748,7 +744,7 @@ def __init__(self, compile_id):
         self.compile_id: Optional[CompileId] = compile_id
         self.attempt = 0
         # Verbose ShapeEnv guards produced.
-        self.shape_env_guards: List[str] = []
+        self.shape_env_guards: list[str] = []
 
     @staticmethod
     def current_compile_id():
@@ -816,7 +812,7 @@ def __init__(self, fake_mode):
         # careful not to accidentally induce guards on the SymInt if
         # you ever do change this in aot_autograd.py; you should check
         # on permutations preferentially.)
-        self.output_strides: Optional[List[Optional[Tuple[int, ...]]]] = None
+        self.output_strides: Optional[list[Optional[tuple[int, ...]]]] = None
         # When this is True, whenever we encounter an int in Dynamo tracing,
         # we will (1) force unspec it and (2) force it as a size-like unbacked
         # integer.  This is currently used when processing certain lists of