[9/N] Apply ruff UP035 rule (#165515)

This is follow-up of #165214 to continue applying ruff UP035 rule to the code base.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/165515
Approved by: https://github.com/Lucaskabela

benchmarks/dynamo/cachebench.py

@@ -6,7 +6,7 @@ import os
 import subprocess
 import sys
 import tempfile
-from typing import Callable
+from collections.abc import Callable
 
 from torch._inductor.utils import fresh_cache
 

benchmarks/dynamo/genai_layers/utils.py

@@ -1,7 +1,8 @@
 import os
 from collections import defaultdict
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Callable, Optional
+from typing import Any, Optional
 
 import matplotlib.pyplot as plt
 

benchmarks/dynamo/torchao_backend.py

@@ -1,4 +1,5 @@
-from typing import Any, Callable
+from collections.abc import Callable
+from typing import Any
 
 import torch
 

benchmarks/functional_autograd_benchmark/functional_autograd_benchmark.py

@@ -1,7 +1,8 @@
 import time
 from argparse import ArgumentParser
 from collections import defaultdict
-from typing import Any, Callable, NamedTuple
+from collections.abc import Callable
+from typing import Any, NamedTuple
 
 import torch
 from torch.autograd import functional

benchmarks/functional_autograd_benchmark/utils.py

@@ -1,5 +1,6 @@
 from collections import defaultdict
-from typing import Callable, Optional, Union
+from collections.abc import Callable
+from typing import Optional, Union
 
 import torch
 from torch import nn, Tensor

benchmarks/gpt_fast/common.py

@@ -1,5 +1,6 @@
 import dataclasses
-from typing import Callable, Optional
+from collections.abc import Callable
+from typing import Optional
 
 
 all_experiments: dict[str, Callable] = {}

benchmarks/inductor_backends/cutlass.py

@@ -9,8 +9,9 @@ import logging
 import time
 from abc import abstractmethod
 from collections import defaultdict
+from collections.abc import Callable
 from dataclasses import asdict, dataclass, field
-from typing import Any, Callable, Optional
+from typing import Any, Optional
 
 from tabulate import tabulate
 from tqdm import tqdm

benchmarks/transformer/attention_bias_benchmarks.py

@@ -1,7 +1,8 @@
 import itertools
+from collections.abc import Callable
 from dataclasses import asdict, dataclass
 from functools import partial
-from typing import Callable, Union
+from typing import Union
 
 import numpy as np
 from tabulate import tabulate

benchmarks/transformer/score_mod.py

@@ -3,10 +3,11 @@ import csv
 import itertools
 import random
 from collections import defaultdict
+from collections.abc import Callable
 from contextlib import nullcontext
 from dataclasses import asdict, dataclass
 from functools import partial
-from typing import Callable, Optional, Union
+from typing import Optional, Union
 
 import numpy as np
 from tabulate import tabulate

benchmarks/transformer/sdpa.py

@@ -1,8 +1,8 @@
 import itertools
 from collections import defaultdict
+from collections.abc import Callable
 from contextlib import nullcontext
 from dataclasses import asdict, dataclass
-from typing import Callable
 
 from tabulate import tabulate
 from tqdm import tqdm

functorch/dim/__init__.py

@@ -3,11 +3,11 @@ from __future__ import annotations
 import dis
 import inspect
 import sys
-from typing import Any, Callable, Optional, TYPE_CHECKING, Union
+from typing import Any, Optional, TYPE_CHECKING, Union
 
 
 if TYPE_CHECKING:
-    from collections.abc import Sequence
+    from collections.abc import Callable, Sequence
 
 import torch
 from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten

functorch/dim/_wrap.py

@@ -5,7 +5,7 @@ Python implementation of function wrapping functionality for functorch.dim.
 from __future__ import annotations
 
 import functools
-from typing import Any, Callable, Optional
+from typing import Any, Optional, TYPE_CHECKING
 
 import torch
 from torch.utils._pytree import tree_map
@@ -15,6 +15,10 @@ from ._enable_all_layers import EnableAllLayers
 from ._tensor_info import TensorInfo
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
 def handle_from_tensor(tensor: torch.Tensor) -> torch.Tensor:
     """Handle tensor conversion for torch function integration."""
     return tensor

functorch/dim/wrap_type.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import functools
+from collections.abc import Callable
 from types import (
     BuiltinMethodType,
     FunctionType,
@@ -12,7 +13,7 @@ from types import (
     MethodDescriptorType,
     WrapperDescriptorType,
 )
-from typing import Any, Callable
+from typing import Any
 
 
 FUNC_TYPES = (

functorch/einops/rearrange.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import functools
-from typing import Callable, TYPE_CHECKING, Union
+from typing import TYPE_CHECKING, Union
 
 import torch
 from functorch.dim import dims  # noqa: F401
@@ -16,7 +16,7 @@ from ._parsing import (
 
 
 if TYPE_CHECKING:
-    from collections.abc import Sequence
+    from collections.abc import Callable, Sequence
 
 __all__ = ["rearrange"]
 

tools/autograd/context.py

@@ -1,5 +1,5 @@
 import functools
-from typing import Callable
+from collections.abc import Callable
 
 from torchgen.api.autograd import NativeFunctionWithDifferentiabilityInfo as NFWDI
 from torchgen.context import native_function_manager

tools/autograd/gen_python_functions.py

@@ -36,7 +36,7 @@ from __future__ import annotations
 import itertools
 import re
 from collections import defaultdict
-from typing import Callable, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 import yaml
 
@@ -77,7 +77,7 @@ from .gen_trace_type import should_trace
 
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable, Sequence
+    from collections.abc import Callable, Iterable, Sequence
 
 
 #

tools/autograd/gen_variable_type.py

@@ -29,7 +29,7 @@
 from __future__ import annotations
 
 import re
-from typing import Callable, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 from torchgen.api import cpp
 from torchgen.api.autograd import (
@@ -106,7 +106,7 @@ from .gen_trace_type import (
 
 
 if TYPE_CHECKING:
-    from collections.abc import Sequence
+    from collections.abc import Callable, Sequence
 
 
 # We don't set or modify grad_fn on these methods. Generally, they return

tools/flight_recorder/components/fr_logger.py

@@ -5,7 +5,8 @@
 # LICENSE file in the root directory of this source tree.
 
 import logging
-from typing import Any, Callable, Optional
+from collections.abc import Callable
+from typing import Any, Optional
 
 
 class FlightRecorderLogger:

tools/github/github_utils.py

@@ -4,12 +4,16 @@ from __future__ import annotations
 
 import json
 import os
-from typing import Any, Callable, cast
+from typing import Any, cast, TYPE_CHECKING
 from urllib.error import HTTPError
 from urllib.parse import quote
 from urllib.request import Request, urlopen
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
 def gh_fetch_url_and_headers(
     url: str,
     *,

tools/linter/adapters/docstring_linter.py

@@ -5,7 +5,7 @@ import json
 import sys
 from functools import cached_property
 from pathlib import Path
-from typing import Any, Callable, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 
 
 _FILE = Path(__file__).absolute()
@@ -18,7 +18,7 @@ else:
     import _linter
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator, Sequence
+    from collections.abc import Callable, Iterator, Sequence
 
 
 GRANDFATHER_LIST = _FILE.parent / "docstring_linter-grandfather.json"

tools/linter/adapters/no_workflows_on_fork.py

@@ -22,11 +22,15 @@ import os
 import re
 from enum import Enum
 from pathlib import Path
-from typing import Any, Callable, NamedTuple, Optional
+from typing import Any, NamedTuple, Optional, TYPE_CHECKING
 
 from yaml import load
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
 # Safely load fast C Yaml loader/dumper if they are available
 try:
     from yaml import CSafeLoader as Loader

tools/nightly.py

@@ -65,10 +65,11 @@ import textwrap
 import time
 import uuid
 from ast import literal_eval
+from collections.abc import Callable
 from datetime import datetime
 from pathlib import Path
 from platform import system as platform_system
-from typing import Any, Callable, cast, NamedTuple, TYPE_CHECKING, TypeVar
+from typing import Any, cast, NamedTuple, TYPE_CHECKING, TypeVar
 
 
 if TYPE_CHECKING:

tools/stats/import_test_stats.py

@@ -7,10 +7,14 @@ import json
 import os
 import shutil
 from pathlib import Path
-from typing import Any, Callable, cast
+from typing import Any, cast, TYPE_CHECKING
 from urllib.request import urlopen
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
 REPO_ROOT = Path(__file__).resolve().parents[2]
 
 

tools/stats/upload_external_contrib_stats.py

@@ -6,13 +6,17 @@ import json
 import os
 import time
 import urllib.parse
-from typing import Any, Callable, cast
+from typing import Any, cast, TYPE_CHECKING
 from urllib.error import HTTPError
 from urllib.request import Request, urlopen
 
 from tools.stats.upload_stats_lib import upload_to_s3
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
 FILTER_OUT_USERS = {
     "pytorchmergebot",
     "facebook-github-bot",

tools/stats/upload_stats_lib.py

@@ -9,12 +9,16 @@ import time
 import zipfile
 from functools import lru_cache
 from pathlib import Path
-from typing import Any, Callable, cast, Optional
+from typing import Any, cast, Optional, TYPE_CHECKING
 
 import boto3  # type: ignore[import]
 import requests
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
 PYTORCH_REPO = "https://api.github.com/repos/pytorch/pytorch"
 
 

tools/testing/target_determination/heuristics/filepath.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from collections import defaultdict
 from functools import lru_cache
 from pathlib import Path
-from typing import Any, Callable
+from typing import Any, TYPE_CHECKING
 from warnings import warn
 
 from tools.testing.target_determination.heuristics.interface import (
@@ -17,6 +17,10 @@ from tools.testing.target_determination.heuristics.utils import (
 from tools.testing.test_run import TestRun
 
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
 REPO_ROOT = Path(__file__).parents[3]
 
 keyword_synonyms: dict[str, list[str]] = {

tools/testing/test_selections.py

@@ -4,7 +4,7 @@ import math
 import os
 import subprocess
 from pathlib import Path
-from typing import Callable, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 from tools.stats.import_test_stats import get_disabled_tests
 from tools.testing.test_run import ShardedTest, TestRun
@@ -19,7 +19,7 @@ except ImportError:
 
 
 if TYPE_CHECKING:
-    from collections.abc import Sequence
+    from collections.abc import Callable, Sequence
 
 
 REPO_ROOT = Path(__file__).resolve().parents[2]