[2/N] Apply ruff UP035 check in torch files (#164054)

This is the result of applying the ruff `UP035` check.
`Callable` is imported from `collections.abc` instead of `typing`.
`TypeAlias` is also imported from `typing`.
This PR is the follow-up of #163947.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/164054
Approved by: https://github.com/ezyang, https://github.com/Skylion007

torch/nn/attention/flex_attention.py

@@ -8,8 +8,9 @@ import itertools
 import math
 import operator
 import warnings
+from collections.abc import Callable
 from enum import Enum
-from typing import Any, Callable, NamedTuple, Optional, Union
+from typing import Any, NamedTuple, Optional, Union
 
 import torch
 from torch import Tensor

torch/nn/common_types.py

@@ -1,5 +1,4 @@
-from typing import Optional, TypeVar, Union
+from typing import Optional, TypeAlias as _TypeAlias, TypeVar, Union
-from typing_extensions import TypeAlias as _TypeAlias
 
 from torch import Tensor
 

torch/nn/functional.py

@@ -3,7 +3,8 @@
 import importlib
 import math
 import warnings
-from typing import Callable, Optional, TYPE_CHECKING, Union
+from collections.abc import Callable
+from typing import Optional, TYPE_CHECKING, Union
 
 import torch
 from torch import _VF, sym_int as _sym_int, Tensor

torch/nn/init.py

@@ -2,7 +2,8 @@
 
 import math
 import warnings
-from typing import Callable, Literal, Optional as _Optional, TypeVar, Union
+from collections.abc import Callable
+from typing import Literal, Optional as _Optional, TypeVar, Union
 from typing_extensions import ParamSpec
 
 import torch

torch/nn/modules/loss.py

@@ -1,5 +1,6 @@
 # mypy: allow-untyped-defs
-from typing import Callable, Optional, Union
+from collections.abc import Callable
+from typing import Optional, Union
 from typing_extensions import deprecated
 
 from torch import Tensor

torch/nn/modules/module.py

@@ -6,8 +6,8 @@ import itertools
 import warnings
 import weakref
 from collections import namedtuple, OrderedDict
-from collections.abc import Iterator, Mapping
+from collections.abc import Callable, Iterator, Mapping
-from typing import Any, Callable, Optional, overload, TypeVar, Union
+from typing import Any, Optional, overload, TypeVar, Union
 from typing_extensions import Self
 
 import torch

torch/nn/modules/transformer.py

@@ -1,7 +1,8 @@
 # mypy: allow-untyped-defs
 import copy
 import warnings
-from typing import Any, Callable, Optional, Union
+from collections.abc import Callable
+from typing import Any, Optional, Union
 
 import torch
 import torch.nn.functional as F

torch/nn/parallel/distributed.py

@@ -9,10 +9,11 @@ import sys
 import warnings
 import weakref
 from collections import defaultdict, deque
+from collections.abc import Callable
 from contextlib import contextmanager
 from dataclasses import dataclass, fields, is_dataclass
 from enum import auto, Enum
-from typing import Any, Callable, Optional, TYPE_CHECKING
+from typing import Any, Optional, TYPE_CHECKING
 
 import torch
 import torch.distributed as dist

torch/nn/utils/_deprecation_utils.py

@@ -1,6 +1,6 @@
 import importlib
 import warnings
-from typing import Callable
+from collections.abc import Callable
 
 
 _MESSAGE_TEMPLATE = (

torch/nn/utils/_expanded_weights/conv_expanded_weights.py

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

torch/nn/utils/_expanded_weights/expanded_weights_impl.py

@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 import functools
+from collections.abc import Callable
 from contextlib import contextmanager
-from typing import Callable
 
 import torch
 from torch._decomp import decomposition_table

torch/nn/utils/clip_grad.py

@@ -4,8 +4,9 @@ import functools
 import types
 import typing
 import warnings
-from typing import Callable, cast, Optional, TypeVar, Union
+from collections.abc import Callable
-from typing_extensions import deprecated, ParamSpec, TypeAlias
+from typing import cast, Optional, TypeAlias, TypeVar, Union
+from typing_extensions import deprecated, ParamSpec
 
 import torch
 from torch import Tensor

torch/nn/utils/rnn.py

@@ -1,6 +1,6 @@
 import warnings
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
-from typing import Any, Callable, NamedTuple, Optional, overload, TypeVar, Union
+from typing import Any, NamedTuple, Optional, overload, TypeVar, Union
 from typing_extensions import Self
 
 import torch

torch/package/analyze/trace_dependencies.py

@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 import sys
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
-from typing import Any, Callable
+from typing import Any
 
 
 __all__ = ["trace_dependencies"]

torch/package/package_exporter.py

@@ -8,12 +8,12 @@ import pickletools
 import platform
 import types
 from collections import defaultdict, OrderedDict
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass
 from enum import Enum
 from importlib.machinery import SourceFileLoader
 from pathlib import Path
-from typing import Any, Callable, cast, IO, Optional, Union
+from typing import Any, cast, IO, Optional, Union
 
 import torch
 from torch.serialization import location_tag, normalize_storage_type

torch/package/package_importer.py

@@ -8,9 +8,9 @@ import linecache
 import os
 import sys
 import types
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from contextlib import contextmanager
-from typing import Any, Callable, cast, Optional, TYPE_CHECKING, Union
+from typing import Any, cast, Optional, TYPE_CHECKING, Union
 from weakref import WeakValueDictionary
 
 import torch

torch/utils/_appending_byte_serializer.py

@@ -1,7 +1,7 @@
 import base64
 import zlib
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
-from typing import Callable, Generic, TypeVar
+from typing import Generic, TypeVar
 
 
 T = TypeVar("T")

torch/utils/_config_module.py

@@ -8,18 +8,10 @@ import os
 import pickle
 import tokenize
 import unittest
+from collections.abc import Callable
 from dataclasses import dataclass
 from types import FunctionType, ModuleType
-from typing import (
+from typing import Any, Generic, NoReturn, Optional, TYPE_CHECKING, TypeVar, Union
-    Any,
-    Callable,
-    Generic,
-    NoReturn,
-    Optional,
-    TYPE_CHECKING,
-    TypeVar,
-    Union,
-)
 from typing_extensions import deprecated
 from unittest import mock
 

torch/utils/_contextlib.py

@@ -6,7 +6,8 @@ import functools
 import inspect
 import sys
 import warnings
-from typing import Any, Callable, cast, TypeVar
+from collections.abc import Callable
+from typing import Any, cast, TypeVar
 
 
 # Used for annotating the decorator usage of _DecoratorContextManager (e.g.,

torch/utils/_cxx_pytree.py

@@ -14,8 +14,8 @@ collection support for PyTorch APIs.
 
 import functools
 import types
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
-from typing import Any, Callable, Optional, overload, TypeVar, Union
+from typing import Any, Optional, overload, TypeVar, Union
 from typing_extensions import deprecated, TypeIs
 
 import torch.utils._pytree as python_pytree

torch/utils/_exposed_in.py

@@ -1,4 +1,5 @@
-from typing import Callable, TypeVar
+from collections.abc import Callable
+from typing import TypeVar
 
 
 F = TypeVar("F")

torch/utils/_foreach_utils.py

@@ -1,5 +1,4 @@
-from typing import Optional
+from typing import Optional, TypeAlias
-from typing_extensions import TypeAlias
 
 import torch
 from torch import Tensor

torch/utils/_functools.py

@@ -1,6 +1,7 @@
 import functools
-from typing import Callable, TypeVar
+from collections.abc import Callable
-from typing_extensions import Concatenate, ParamSpec
+from typing import Concatenate, TypeVar
+from typing_extensions import ParamSpec
 
 
 _P = ParamSpec("_P")

torch/utils/_pytree.py

@@ -24,11 +24,10 @@ import threading
 import types
 import warnings
 from collections import defaultdict, deque, namedtuple, OrderedDict
-from collections.abc import Hashable, Iterable, Mapping, Sequence
+from collections.abc import Callable, Hashable, Iterable, Mapping, Sequence
 from enum import Enum
 from typing import (
     Any,
-    Callable,
     cast,
     ClassVar,
     Final,

torch/utils/_stats.py

@@ -4,7 +4,8 @@
 import collections
 import functools
 from collections import OrderedDict
-from typing import Callable, TypeVar
+from collections.abc import Callable
+from typing import TypeVar
 from typing_extensions import ParamSpec
 
 

torch/utils/_strobelight/cli_function_profiler.py

@@ -6,9 +6,9 @@ import os
 import re
 import subprocess
 import time
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from threading import Lock
-from typing import Any, Callable, Optional, TypeVar
+from typing import Any, Optional, TypeVar
 from typing_extensions import ParamSpec
 
 

torch/utils/_sympy/functions.py

@@ -3,7 +3,8 @@ import functools
 import math
 import operator
 import sys
-from typing import Callable, Optional, SupportsFloat, TYPE_CHECKING, TypeVar, Union
+from collections.abc import Callable
+from typing import Optional, SupportsFloat, TYPE_CHECKING, TypeVar, Union
 from typing_extensions import TypeVarTuple, Unpack
 
 import sympy

torch/utils/_sympy/value_ranges.py

@@ -7,17 +7,17 @@ import itertools
 import logging
 import math
 import operator
+from collections.abc import Callable
 from typing import (
-    Callable,
     Generic,
     Optional,
     overload,
     SupportsFloat,
     TYPE_CHECKING,
+    TypeGuard,
     TypeVar,
     Union,
 )
-from typing_extensions import TypeGuard
 
 import sympy
 from sympy.logic.boolalg import Boolean as SympyBoolean, BooleanAtom

torch/utils/_thunk.py

@@ -1,4 +1,5 @@
-from typing import Callable, Generic, Optional, TypeVar
+from collections.abc import Callable
+from typing import Generic, Optional, TypeVar
 
 
 R = TypeVar("R")

torch/utils/benchmark/utils/_stubs.py

@@ -1,4 +1,5 @@
-from typing import Any, Callable
+from typing import Any
+from collections.abc import Callable
 from typing_extensions import Protocol, runtime_checkable
 
 

torch/utils/benchmark/utils/compile.py

@@ -1,5 +1,6 @@
 # mypy: allow-untyped-defs
-from typing import Any, Callable, cast, Optional, Union
+from typing import Any, cast, Optional, Union
+from collections.abc import Callable
 
 import torch
 import torch._dynamo

torch/utils/benchmark/utils/fuzzer.py

@@ -1,7 +1,8 @@
 # mypy: allow-untyped-defs
 import functools
 import itertools as it
-from typing import Any, Callable, Optional, Union
+from typing import Any, Optional, Union
+from collections.abc import Callable
 
 import torch
 

torch/utils/benchmark/utils/timer.py

@@ -2,7 +2,8 @@
 import enum
 import timeit
 import textwrap
-from typing import overload, Any, Callable, NoReturn, Optional, Union
+from typing import overload, Any, NoReturn, Optional, Union
+from collections.abc import Callable
 
 import torch
 from torch.utils.benchmark.utils import common, cpp_jit

torch/utils/benchmark/utils/valgrind_wrapper/timer_interface.py

@@ -11,8 +11,9 @@ import subprocess
 import sys
 import textwrap
 from typing import (
-    cast, Any, Callable, NamedTuple,
+    cast, Any, NamedTuple,
     Optional, Union, TYPE_CHECKING)
+from collections.abc import Callable
 from collections.abc import Iterator
 
 import torch

torch/utils/bundled_inputs.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 # mypy: allow-untyped-defs
-from typing import Any, TypeVar, Optional, NamedTuple, Union, Callable
+from typing import Any, TypeVar, Optional, NamedTuple, Union
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 import textwrap
 import torch
 from torch._C import TupleType, ListType

torch/utils/data/_utils/collate.py

@@ -12,7 +12,8 @@ import collections
 import contextlib
 import copy
 import re
-from typing import Callable, Optional, Union
+from collections.abc import Callable
+from typing import Optional, Union
 
 import torch
 

torch/utils/data/dataloader.py

@@ -16,7 +16,8 @@ import os
 import queue
 import threading
 import warnings
-from typing import Any, Callable, Generic, Optional, TYPE_CHECKING, TypeVar, Union
+from collections.abc import Callable
+from typing import Any, Generic, Optional, TYPE_CHECKING, TypeVar, Union
 from typing_extensions import Self
 
 import torch

torch/utils/data/datapipes/datapipe.py

@@ -1,7 +1,7 @@
 import functools
 import pickle
-from collections.abc import Iterable, Iterator
+from collections.abc import Callable, Iterable, Iterator
-from typing import Callable, Optional, TypeVar
+from typing import Optional, TypeVar
 
 from torch.utils._import_utils import import_dill
 from torch.utils.data.datapipes._hook_iterator import _SnapshotState

torch/utils/data/datapipes/iter/callable.py

@@ -1,8 +1,8 @@
 # mypy: allow-untyped-defs
 import functools
 from collections import namedtuple
-from collections.abc import Iterator, Sized
+from collections.abc import Callable, Iterator, Sized
-from typing import Any, Callable, Optional, TypeVar, Union
+from typing import Any, Optional, TypeVar, Union
 
 import torch
 from torch.utils.data._utils.collate import default_collate

torch/utils/data/datapipes/iter/combining.py

@@ -3,8 +3,8 @@ import copy as copymodule
 import warnings
 from abc import ABC, abstractmethod
 from collections import deque
-from collections.abc import Iterator, Sized
+from collections.abc import Callable, Iterator, Sized
-from typing import Any, Callable, Literal, Optional, TypeVar
+from typing import Any, Literal, Optional, TypeVar
 
 from torch.utils.data.datapipes._decorator import functional_datapipe
 from torch.utils.data.datapipes._hook_iterator import _SnapshotState

torch/utils/data/datapipes/iter/grouping.py

@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 from collections import defaultdict
-from collections.abc import Iterator, Sized
+from collections.abc import Callable, Iterator, Sized
-from typing import Any, Callable, Optional, TypeVar
+from typing import Any, Optional, TypeVar
 
 from torch.utils.data.datapipes._decorator import functional_datapipe
 from torch.utils.data.datapipes.datapipe import DataChunk, IterDataPipe

torch/utils/data/datapipes/iter/routeddecoder.py

@@ -1,6 +1,6 @@
-from collections.abc import Iterable, Iterator, Sized
+from collections.abc import Callable, Iterable, Iterator, Sized
 from io import BufferedIOBase
-from typing import Any, Callable
+from typing import Any
 
 from torch.utils.data.datapipes._decorator import functional_datapipe
 from torch.utils.data.datapipes.datapipe import IterDataPipe

torch/utils/data/datapipes/iter/selecting.py

@@ -1,6 +1,6 @@
 # mypy: allow-untyped-defs
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
-from typing import Callable, TypeVar
+from typing import TypeVar
 
 from torch.utils.data.datapipes._decorator import functional_datapipe
 from torch.utils.data.datapipes.dataframe import dataframe_wrapper as df_wrapper

torch/utils/data/datapipes/map/callable.py

@@ -1,5 +1,6 @@
 # mypy: allow-untyped-defs
-from typing import Callable, TypeVar
+from collections.abc import Callable
+from typing import TypeVar
 
 from torch.utils.data.datapipes._decorator import functional_datapipe
 from torch.utils.data.datapipes.datapipe import MapDataPipe

torch/utils/data/datapipes/utils/common.py

@@ -4,9 +4,9 @@ import functools
 import inspect
 import os
 import warnings
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from io import IOBase
-from typing import Any, Callable, Optional, Union
+from typing import Any, Optional, Union
 
 from torch.utils._import_utils import dill_available
 

torch/utils/flop_counter.py

@@ -2,7 +2,8 @@
 import torch
 from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten
 from .module_tracker import ModuleTracker
-from typing import Any, Optional, Union, TypeVar, Callable
+from typing import Any, Optional, Union, TypeVar
+from collections.abc import Callable
 from collections.abc import Iterator
 from typing_extensions import ParamSpec
 from collections import defaultdict

torch/xpu/__init__.py

@@ -9,8 +9,9 @@ This package is lazily initialized, so you can always import it, and use
 
 import threading
 import traceback
+from collections.abc import Callable
 from functools import lru_cache
-from typing import Any, Callable, Optional, Union
+from typing import Any, Optional, Union
 
 import torch
 import torch._C

torch/xpu/_gpu_trace.py

@@ -1,4 +1,4 @@
-from typing import Callable
+from collections.abc import Callable
 
 from torch._utils import CallbackRegistry