PEP585: More UP006 fixes (pytorch#146392)

This should be the final PR before we can enable RUFF UP006.

Pull Request resolved: pytorch#146392
Approved by: https://github.com/justinchuby, https://github.com/albanD, https://github.com/Skylion007

Author: aorenste

test/distributed/_composable/fsdp/test_fully_shard_ignore_params.py

@@ -1,7 +1,6 @@
 # Owner(s): ["oncall: distributed"]
 
 import sys
-from typing import List
 
 import torch
 import torch.distributed as dist
@@ -119,7 +118,7 @@ def _find_name_param_mappings(module: torch.nn.Module, prefix: str):
 
 
 def _discover_ddp_ignored_params(module: torch.nn.Module, prefix: str):
-    ddp_ignore_parameters: List[str] = []
+    ddp_ignore_parameters: list[str] = []
     if isinstance(module, FSDP2):
         ddp_ignore_parameters = [name for name, _ in module.named_parameters(prefix)]
     else:

test/distributed/test_c10d_functional_native.py

@@ -3,7 +3,7 @@
 import threading
 import unittest
 from datetime import timedelta
-from typing import List, Optional
+from typing import Optional
 
 import torch
 import torch.distributed as dist
@@ -576,24 +576,24 @@ def __init__(self) -> None:
         self.waits = 0
         self.dels = 0
 
-    def broadcast(self, tensor_list: List[torch.Tensor], opts: object) -> dist.Work:
+    def broadcast(self, tensor_list: list[torch.Tensor], opts: object) -> dist.Work:
         return _DummyWork(self)
 
     def allgather_into_tensor_coalesced(
         self,
-        output_lists: List[torch.Tensor],
-        input_list: List[torch.Tensor],
+        output_lists: list[torch.Tensor],
+        input_list: list[torch.Tensor],
         opts: object,
     ) -> dist.Work:
         return _DummyWork(self)
 
-    def allreduce(self, tensors: List[torch.Tensor], opts: object) -> dist.Work:
+    def allreduce(self, tensors: list[torch.Tensor], opts: object) -> dist.Work:
         return _DummyWork(self)
 
     def reduce_scatter_tensor_coalesced(
         self,
-        outputTensors: List[torch.Tensor],
-        inputTensors: List[torch.Tensor],
+        outputTensors: list[torch.Tensor],
+        inputTensors: list[torch.Tensor],
         opts: object,
     ) -> dist.Work:
         return _DummyWork(self)

test/functorch/test_ac_logging.py

@@ -1,5 +1,4 @@
 # Owner(s): ["module: functorch"]
-from typing import Dict, List, Tuple
 from unittest.mock import MagicMock, patch
 
 from torch._functorch._activation_checkpointing.ac_logging_utils import (
@@ -33,17 +32,17 @@ def setUp(self) -> None:
 
         self.graph.nodes = [self.node1, self.node2]
 
-        self.all_recomputable_banned_nodes: List[Node] = [self.node1]
-        self.saved_node_idxs: List[int] = [0]
-        self.recomputable_node_idxs: List[int] = []
+        self.all_recomputable_banned_nodes: list[Node] = [self.node1]
+        self.saved_node_idxs: list[int] = [0]
+        self.recomputable_node_idxs: list[int] = []
         self.expected_runtime: int = 100
-        self.memories_banned_nodes: List[int] = [50]
-        self.runtimes_banned_nodes: List[int] = [10]
-        self.min_cut_saved_values: List[Node] = [self.node1]
+        self.memories_banned_nodes: list[int] = [50]
+        self.runtimes_banned_nodes: list[int] = [10]
+        self.min_cut_saved_values: list[Node] = [self.node1]
 
     def test_create_joint_graph_node_information(self) -> None:
-        recomputable_node_info: Dict[str, int] = {"node1": 0}
-        expected_output: Dict[str, Dict] = {
+        recomputable_node_info: dict[str, int] = {"node1": 0}
+        expected_output: dict[str, dict] = {
             "node1": {
                 "index": 0,
                 "name": "node1",
@@ -68,12 +67,12 @@ def test_create_joint_graph_node_information(self) -> None:
         self.assertEqual(result, expected_output)
 
     def test_create_joint_graph_edges(self) -> None:
-        expected_edges: List[Tuple[str, str]] = [("node1", "node2")]
+        expected_edges: list[tuple[str, str]] = [("node1", "node2")]
         result = create_joint_graph_edges(self.graph)
         self.assertEqual(result, expected_edges)
 
     def test_create_activation_checkpointing_logging_structure_payload(self) -> None:
-        input_joint_graph_node_information: Dict[str, Dict] = {
+        input_joint_graph_node_information: dict[str, dict] = {
             "node1": {
                 "index": 0,
                 "name": "node1",
@@ -85,8 +84,8 @@ def test_create_activation_checkpointing_logging_structure_payload(self) -> None
                 "recomputable_candidate_info": {"recomputable_node_idx": 0},
             }
         }
-        joint_graph_edges: List[Tuple[str, str]] = [("node1", "node2")]
-        expected_payload: Dict[str, any] = {
+        joint_graph_edges: list[tuple[str, str]] = [("node1", "node2")]
+        expected_payload: dict[str, any] = {
             "Joint Graph Size": 2,
             "Joint Graph Edges": {"Total": 1, "Edges": joint_graph_edges},
             "Joint Graph Node Information": input_joint_graph_node_information,

test/onnx/test_onnxscript_runtime.py

@@ -2,7 +2,7 @@
 
 """Test the support on onnxscript in PyTorch-ONNX converter with onnxruntime."""
 
-from typing import Sequence
+from typing import Sequence  # noqa: UP035
 
 import onnx_test_common
 import onnxscript

test/onnx/torchlib/error_reproduction.py

@@ -9,7 +9,7 @@
 import sys
 import time
 import traceback
-from typing import Any, Mapping
+from typing import Any, TYPE_CHECKING
 
 import numpy as np
 
@@ -20,6 +20,10 @@
 import torch
 
 
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+
 _REPRODUCTION_TEMPLATE = '''\
 import google.protobuf.text_format
 import numpy as np

torch/_C/_monitor.pyi

@@ -3,7 +3,7 @@
 import datetime
 from enum import Enum
 from types import TracebackType
-from typing import Callable, Optional, Type
+from typing import Callable, Optional
 
 class Aggregation(Enum):
     VALUE = ...
@@ -48,7 +48,7 @@ class _WaitCounterTracker:
     def __enter__(self) -> None: ...
     def __exit__(
         self,
-        exec_type: Optional[Type[BaseException]] = None,
+        exec_type: Optional[type[BaseException]] = None,
         exec_value: Optional[BaseException] = None,
         traceback: Optional[TracebackType] = None,
     ) -> None: ...

torch/_decomp/__init__.py

@@ -4,17 +4,7 @@
 from collections.abc import Sequence
 from functools import lru_cache, partial, wraps
 from itertools import chain
-from typing import (
-    Callable,
-    Dict,
-    FrozenSet,
-    List,
-    Optional,
-    Set,
-    TYPE_CHECKING,
-    TypeVar,
-    Union,
-)
+from typing import Callable, Optional, TYPE_CHECKING, TypeVar, Union
 from typing_extensions import ParamSpec
 
 

torch/_dynamo/polyfills/__init__.py

@@ -11,7 +11,7 @@
 import types
 from collections.abc import MutableMapping, Sequence
 from itertools import repeat as _repeat
-from typing import Any, Callable, List, TYPE_CHECKING
+from typing import Any, Callable, TYPE_CHECKING
 
 import torch
 

torch/_dynamo/utils.py

@@ -59,7 +59,6 @@
     Generic,
     Optional,
     overload,
-    Set,
     TypeVar,
     Union,
 )
@@ -1393,7 +1392,7 @@ def default(self, o):
             except Exception:
                 return "Value is not JSON serializable"
 
-    keys_to_scrub: Set[Any] = set()
+    keys_to_scrub: set[Any] = set()
     inductor_conf_str = None
     inductor_config_copy = (
         torch._inductor.config.get_config_copy() if torch._inductor.config else None

torch/_dynamo/variables/base.py

@@ -16,8 +16,9 @@
 """
 
 import collections
+from collections.abc import Sequence
 from enum import Enum
-from typing import Any, Callable, Optional, Sequence, TYPE_CHECKING
+from typing import Any, Callable, Optional, TYPE_CHECKING
 
 from .. import variables
 from ..current_scope_id import current_scope_id

torch/_dynamo/variables/builtin.py

@@ -11,8 +11,8 @@
 import types
 import typing
 from collections import defaultdict, OrderedDict
-from collections.abc import KeysView
-from typing import Callable, Sequence, TYPE_CHECKING, Union
+from collections.abc import KeysView, Sequence
+from typing import Callable, TYPE_CHECKING, Union
 
 import torch
 from torch import sym_float, sym_int

torch/_dynamo/variables/functions.py

@@ -30,7 +30,7 @@
 import sys
 import types
 from collections.abc import Sequence
-from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING, TypeVar
+from typing import Any, Callable, Optional, TYPE_CHECKING, TypeVar
 from typing_extensions import Never
 from unittest.mock import patch
 
@@ -517,7 +517,7 @@ def has_unpack_var_sequence(self, tx):
     def has_force_unpack_var_sequence(self, tx) -> builtins.bool:
         return True
 
-    def force_unpack_var_sequence(self, tx) -> List[VariableTracker]:
+    def force_unpack_var_sequence(self, tx) -> list[VariableTracker]:
         result = []
         while True:
             try:
@@ -547,8 +547,8 @@ def call_method(
         self,
         tx: "InstructionTranslator",
         name: str,
-        args: "List[VariableTracker]",
-        kwargs: "Dict[str, VariableTracker]",
+        args: "list[VariableTracker]",
+        kwargs: "dict[str, VariableTracker]",
     ) -> "VariableTracker":
         if name == "__next__":
             return self.next_variable(tx)

torch/_export/__init__.py

@@ -16,7 +16,7 @@
 from contextlib import contextmanager
 from functools import lru_cache
 
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Optional, Union
 from unittest.mock import patch
 
 import torch

torch/_functorch/_activation_checkpointing/ac_logging_utils.py

@@ -1,6 +1,6 @@
 import json
 import logging
-from typing import Any, Dict, List, Tuple
+from typing import Any
 
 from torch._logging import trace_structured
 from torch.fx import Graph, Node
@@ -11,9 +11,9 @@
 
 def create_joint_graph_node_information(
     joint_graph: Graph,
-    recomputable_node_info: Dict[str, int],
-) -> Dict[str, Any]:
-    joint_graph_node_information: Dict[str, Any] = {}
+    recomputable_node_info: dict[str, int],
+) -> dict[str, Any]:
+    joint_graph_node_information: dict[str, Any] = {}
 
     for i, joint_graph_node in enumerate(joint_graph.nodes):
         is_recomputable_candidate: bool = (
@@ -22,7 +22,7 @@ def create_joint_graph_node_information(
         tensor_meta = joint_graph_node.meta.get("tensor_meta")
         shape = getattr(tensor_meta, "shape", []) if tensor_meta else []
 
-        node_info: Dict[str, Any] = {
+        node_info: dict[str, Any] = {
             "index": i,
             "name": joint_graph_node.name,
             "is_recomputable_candidate": is_recomputable_candidate,
@@ -43,8 +43,8 @@ def create_joint_graph_node_information(
     return joint_graph_node_information
 
 
-def create_joint_graph_edges(joint_graph: Graph) -> List[Tuple[str, str]]:
-    joint_graph_edges: List[Tuple[str, str]] = [
+def create_joint_graph_edges(joint_graph: Graph) -> list[tuple[str, str]]:
+    joint_graph_edges: list[tuple[str, str]] = [
         (inp.name, node.name)
         for node in joint_graph.nodes
         for inp in node.all_input_nodes
@@ -54,17 +54,17 @@ def create_joint_graph_edges(joint_graph: Graph) -> List[Tuple[str, str]]:
 
 def create_activation_checkpointing_logging_structure_payload(
     joint_graph: Graph,
-    joint_graph_node_information: Dict[str, Any],
-    joint_graph_edges: List[Tuple[str, str]],
-    all_recomputable_banned_nodes: List[Node],
+    joint_graph_node_information: dict[str, Any],
+    joint_graph_edges: list[tuple[str, str]],
+    all_recomputable_banned_nodes: list[Node],
     expected_runtime: float,
-    saved_node_idxs: List[int],
-    recomputable_node_idxs: List[int],
-    memories_banned_nodes: List[float],
-    runtimes_banned_nodes: List[float],
-    min_cut_saved_values: List[Node],
-) -> Dict[str, Any]:
-    activation_checkpointing_logging_structure_payload: Dict[str, Any] = {
+    saved_node_idxs: list[int],
+    recomputable_node_idxs: list[int],
+    memories_banned_nodes: list[float],
+    runtimes_banned_nodes: list[float],
+    min_cut_saved_values: list[Node],
+) -> dict[str, Any]:
+    activation_checkpointing_logging_structure_payload: dict[str, Any] = {
         "Joint Graph Size": len(joint_graph.nodes),
         "Joint Graph Edges": {
             "Total": len(joint_graph_edges),
@@ -86,15 +86,15 @@ def create_activation_checkpointing_logging_structure_payload(
 
 def create_structured_trace_for_min_cut_info(
     joint_graph: Graph,
-    all_recomputable_banned_nodes: List[Node],
-    saved_node_idxs: List[int],
-    recomputable_node_idxs: List[int],
+    all_recomputable_banned_nodes: list[Node],
+    saved_node_idxs: list[int],
+    recomputable_node_idxs: list[int],
     expected_runtime: float,
-    memories_banned_nodes: List[float],
-    runtimes_banned_nodes: List[float],
-    min_cut_saved_values: List[Node],
+    memories_banned_nodes: list[float],
+    runtimes_banned_nodes: list[float],
+    min_cut_saved_values: list[Node],
 ) -> None:
-    recomputable_node_info: Dict[str, int] = {
+    recomputable_node_info: dict[str, int] = {
         node.name: idx for idx, node in enumerate(all_recomputable_banned_nodes)
     }
     joint_graph_node_information = create_joint_graph_node_information(

torch/_functorch/_aot_autograd/subclass_parametrization.py

@@ -1,6 +1,7 @@
 import dataclasses
 import itertools
-from typing import Any, Iterable, Union
+from collections.abc import Iterable
+from typing import Any, Union
 
 import torch
 from torch.utils._python_dispatch import is_traceable_wrapper_subclass