[8/N] Apply ruff UP035 rule (#165214)

This is follow-up of #164653 to continue applying `UP035` fixes. The purpose is to finally enable this rule.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/165214
Approved by: https://github.com/ezyang

test/fx/test_matcher_utils.py

@@ -2,7 +2,7 @@
 
 import os
 import sys
-from typing import Callable
+from collections.abc import Callable
 
 import torch
 import torch.nn.functional as F

test/test_fx.py

@@ -35,7 +35,8 @@ from torch.fx.experimental.rewriter import RewritingTracer
 from torch.fx.operator_schemas import get_signature_for_torch_op
 from copy import deepcopy
 from collections import namedtuple
-from typing import Any, Callable, NamedTuple, Optional, Union
+from typing import Any, NamedTuple, Optional, Union
+from collections.abc import Callable
 
 import torch
 

test/test_fx_experimental.py

@@ -12,7 +12,8 @@ import tempfile
 import typing
 import unittest
 from types import BuiltinFunctionType
-from typing import Callable, NamedTuple, Optional, Union
+from typing import NamedTuple, Optional, Union
+from collections.abc import Callable
 
 import torch
 import torch.fx.experimental.meta_tracer

torch/distributed/pipelining/microbatch.py

@@ -2,7 +2,8 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates
 import logging
 import operator
-from typing import Any, Optional, Sequence
+from collections.abc import Sequence
+from typing import Any, Optional
 
 import torch
 from torch.fx.node import map_aggregate

torch/fx/_compatibility.py

@@ -1,5 +1,6 @@
 import textwrap
-from typing import Any, Callable, TypeVar
+from collections.abc import Callable
+from typing import Any, TypeVar
 
 
 _BACK_COMPAT_OBJECTS: dict[Any, None] = {}

torch/fx/_graph_pickler.py

@@ -3,7 +3,8 @@ import importlib
 import io
 import pickle
 from abc import abstractmethod
-from typing import Any, Callable, NewType, Optional, TypeVar, Union
+from collections.abc import Callable
+from typing import Any, NewType, Optional, TypeVar, Union
 from typing_extensions import override, Self
 
 import torch

torch/fx/_pytree.py

@@ -1,5 +1,6 @@
 from collections import namedtuple
-from typing import Any, Callable, Optional, TypeVar
+from collections.abc import Callable
+from typing import Any, Optional, TypeVar
 from typing_extensions import NamedTuple
 
 import torch.return_types

torch/fx/_symbolic_trace.py

@@ -9,10 +9,10 @@ import logging
 import math
 import os
 import warnings
+from collections.abc import Callable
 from itertools import chain
 from types import CodeType, FunctionType, ModuleType
-from typing import Any, Callable, get_args, NamedTuple, Optional, Union
+from typing import Any, get_args, NamedTuple, Optional, TypeAlias, Union
-from typing_extensions import TypeAlias
 
 import torch
 import torch.utils._pytree as pytree

torch/fx/experimental/_dynamism.py

@@ -1,5 +1,6 @@
 import re
-from typing import Any, Callable, Union
+from collections.abc import Callable
+from typing import Any, Union
 
 import torch
 from torch.utils._pytree import tree_flatten_with_path, tree_map

torch/fx/experimental/const_fold.py

@@ -1,6 +1,7 @@
 # mypy: allow-untyped-defs
 import re
-from typing import Callable, Optional, Union
+from collections.abc import Callable
+from typing import Optional, Union
 
 import torch.fx
 from torch.fx.node import map_arg

torch/fx/experimental/graph_gradual_typechecker.py

@@ -1,8 +1,9 @@
 # mypy: allow-untyped-defs
 import itertools
 import operator
+from collections.abc import Callable
 from functools import reduce
-from typing import Callable, TypeVar
+from typing import TypeVar
 from typing_extensions import ParamSpec
 
 import sympy

torch/fx/experimental/meta_tracer.py

@@ -2,7 +2,8 @@
 import builtins
 import functools
 import warnings
-from typing import Any, Callable, Optional, Union
+from collections.abc import Callable
+from typing import Any, Optional, Union
 
 import torch
 import torch.fx

torch/fx/experimental/migrate_gradual_types/constraint_generator.py

@@ -1,8 +1,8 @@
 # mypy: allow-untyped-defs
 import operator
 import warnings
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
-from typing import Callable, TypeVar
+from typing import TypeVar
 from typing_extensions import ParamSpec
 
 import torch

torch/fx/experimental/migrate_gradual_types/constraint_transformation.py

@@ -1,7 +1,7 @@
 # mypy: ignore-errors
 import copy
 import itertools
-from typing import Callable
+from collections.abc import Callable
 
 from torch.fx.experimental.migrate_gradual_types.constraint import (
     ApplyBroadcasting,

torch/fx/experimental/normalize.py

@@ -1,6 +1,7 @@
 # mypy: allow-untyped-defs
 import operator
-from typing import Any, Callable, Optional
+from collections.abc import Callable
+from typing import Any, Optional
 
 import torch
 import torch.fx

torch/fx/experimental/proxy_tensor.py

@@ -16,12 +16,12 @@ import typing
 import typing_extensions
 import weakref
 from collections import defaultdict, OrderedDict
-from collections.abc import Generator, Mapping, Sequence
+from collections.abc import Callable, Generator, Mapping, Sequence
 from contextlib import _GeneratorContextManager, contextmanager, ExitStack, nullcontext
 from dataclasses import dataclass
 from typing import (
     Any,
-    Callable,
+    Concatenate,
     Optional,
     overload,
     Protocol,
@@ -29,7 +29,7 @@ from typing import (
     TypeVar,
     Union,
 )
-from typing_extensions import Concatenate, ParamSpec, Self, TypeVarTuple, Unpack
+from typing_extensions import ParamSpec, Self, TypeVarTuple, Unpack
 from weakref import WeakKeyDictionary
 
 import torch

torch/fx/experimental/recording.py

@@ -3,8 +3,9 @@ import functools
 import inspect
 import itertools
 import logging
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Callable, Optional, Union
+from typing import Any, Optional, Union
 
 import torch
 import torch.utils._pytree as pytree

torch/fx/experimental/rewriter.py

@@ -5,8 +5,9 @@ import copy
 import functools
 import inspect
 import textwrap
+from collections.abc import Callable
 from types import FunctionType
-from typing import Any, Callable, cast, Optional, Union
+from typing import Any, cast, Optional, Union
 
 import torch
 from torch._sources import normalize_source_lines

torch/fx/experimental/symbolic_shapes.py

@@ -31,23 +31,24 @@ import sys
 import threading
 import traceback
 from collections import Counter, defaultdict
-from collections.abc import Generator, Iterator, Mapping, Sequence
+from collections.abc import Callable, Generator, Iterator, Mapping, Sequence
 from contextlib import _GeneratorContextManager, contextmanager
 from dataclasses import asdict, dataclass, field
 from enum import Enum
 from typing import (
     Any,
-    Callable,
     cast,
     Generic,
     NamedTuple,
     NoReturn,
     Optional,
     TYPE_CHECKING,
+    TypeAlias,
+    TypeGuard,
     TypeVar,
     Union,
 )
-from typing_extensions import deprecated, ParamSpec, TypeAlias, TypeGuard
+from typing_extensions import deprecated, ParamSpec
 
 import torch
 import torch.fx

torch/fx/experimental/unification/multipledispatch/core.py

@@ -1,6 +1,7 @@
 # mypy: allow-untyped-defs
 import inspect
-from typing import Any, Callable, TypeVar
+from collections.abc import Callable
+from typing import Any, TypeVar
 from typing_extensions import TypeVarTuple, Unpack
 
 from .dispatcher import Dispatcher, MethodDispatcher

torch/fx/experimental/validator.py

@@ -4,8 +4,9 @@ import functools
 import logging
 import math
 import operator
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Callable, Optional, Union
+from typing import Any, Optional, Union
 
 import sympy
 

torch/fx/graph.py

@@ -12,10 +12,10 @@ import re
 import typing
 import warnings
 from collections import defaultdict
-from collections.abc import Iterable, Iterator
+from collections.abc import Callable, Iterable, Iterator
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any, Callable, Literal, NamedTuple, Optional, TYPE_CHECKING
+from typing import Any, Literal, NamedTuple, Optional, TYPE_CHECKING
 
 import torch
 import torch.utils._pytree as pytree

torch/fx/graph_module.py

@@ -7,8 +7,9 @@ import os
 import sys
 import traceback
 import warnings
+from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Callable, Optional, Union
+from typing import Any, Optional, Union
 
 import torch
 import torch.nn as nn

torch/fx/node.py

@@ -4,9 +4,9 @@ import inspect
 import logging
 import operator
 import types
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Callable, Iterable, Mapping, Sequence
-from typing import Any, Callable, Optional, TYPE_CHECKING, Union
+from typing import Any, Optional, TYPE_CHECKING, TypeAlias, Union
-from typing_extensions import ParamSpec, TypeAlias, TypeVar
+from typing_extensions import ParamSpec, TypeVar
 
 import torch
 from torch._C import _fx_map_aggregate, _fx_map_arg, _NodeBase

torch/fx/operator_schemas.py

@@ -5,7 +5,8 @@ import numbers
 import types
 import typing
 import warnings
-from typing import Any, Callable, cast, NamedTuple, Optional, TYPE_CHECKING
+from collections.abc import Callable
+from typing import Any, cast, NamedTuple, Optional, TYPE_CHECKING
 
 import torch
 from torch._jit_internal import boolean_dispatched

torch/fx/passes/graph_transform_observer.py

@@ -1,6 +1,7 @@
 # mypy: allow-untyped-defs
 import os
-from typing import Callable, Optional, TypeVar
+from collections.abc import Callable
+from typing import Optional, TypeVar
 
 from torch.fx import Graph, Node
 from torch.fx._compatibility import compatibility

torch/fx/passes/infra/pass_manager.py

@@ -1,9 +1,9 @@
 # mypy: allow-untyped-defs
 import inspect
 import logging
+from collections.abc import Callable
 from functools import wraps
 from queue import Queue
-from typing import Callable
 
 import torch.nn as nn
 from torch.fx._compatibility import compatibility

torch/fx/passes/net_min_base.py

@@ -1,7 +1,8 @@
 # mypy: allow-untyped-defs
 import logging
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Callable, cast, Optional
+from typing import Any, cast, Optional
 
 import torch
 import torch.fx

torch/fx/passes/param_fetch.py

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

torch/fx/passes/pass_manager.py

@@ -1,8 +1,9 @@
 # mypy: allow-untyped-defs
 import logging
+from collections.abc import Callable
 from functools import wraps
 from inspect import unwrap
-from typing import Callable, Optional
+from typing import Optional
 
 
 logger = logging.getLogger(__name__)

torch/fx/passes/reinplace.py

@@ -2,8 +2,9 @@
 import _operator
 import itertools
 from collections import defaultdict
+from collections.abc import Callable
 from enum import Enum
-from typing import Any, Callable
+from typing import Any
 
 import torch
 from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode

torch/fx/passes/split_module.py

@@ -2,7 +2,8 @@
 import inspect
 import logging
 from collections import OrderedDict
-from typing import Any, Callable, Optional
+from collections.abc import Callable
+from typing import Any, Optional
 
 import torch
 from torch.fx._compatibility import compatibility

torch/fx/passes/utils/source_matcher_utils.py

@@ -1,7 +1,8 @@
 import logging
 import os
+from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import Any, Callable, Optional
+from typing import Any, Optional
 
 from torch.fx._compatibility import compatibility
 from torch.fx.graph import Graph

torch/fx/proxy.py

@@ -10,9 +10,9 @@ import operator
 import sys
 import traceback
 from collections import OrderedDict
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
 from dataclasses import fields, is_dataclass
-from typing import Any, Callable, Optional
+from typing import Any, Optional
 
 import torch
 import torch.fx.traceback as fx_traceback

torch/fx/subgraph_rewriter.py

@@ -1,6 +1,7 @@
 import copy
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any, Callable, NamedTuple, Optional, TYPE_CHECKING, Union
+from typing import Any, NamedTuple, Optional, TYPE_CHECKING, Union
 
 import torch