[1/N] Use TYPE_CHECKING (pytorch#165852)

This PR moves typing imports into the `TYPE_CHECKING` block.

Pull Request resolved: pytorch#165852
Approved by: https://github.com/Lucaskabela

Author: cyyever

torch/_export/utils.py

@@ -24,6 +24,8 @@
 
 
 if TYPE_CHECKING:
+    import sympy
+
     from torch._export.passes.lift_constants_pass import ConstantAttrMap
     from torch._ops import OperatorBase
     from torch.export import ExportedProgram
@@ -433,8 +435,6 @@ def _check_symint(
 def _check_input_constraints_for_graph(
     input_placeholders: list[torch.fx.Node], flat_args_with_path, range_constraints
 ) -> None:
-    import sympy  # noqa: TC002
-
     if len(flat_args_with_path) != len(input_placeholders):
         raise RuntimeError(
             "Unexpected number of inputs "

torch/ao/nn/qat/dynamic/modules/linear.py

@@ -4,7 +4,7 @@
 
 
 if TYPE_CHECKING:
-    from torch.ao.quantization.qconfig import QConfig  # noqa: TC004
+    from torch.ao.quantization.qconfig import QConfig
 
 
 __all__ = ["Linear"]

torch/distributed/checkpoint/state_dict_saver.py

@@ -6,15 +6,12 @@
 from concurrent.futures import Future
 from dataclasses import dataclass
 from enum import Enum
-from typing import cast, Optional, Union
+from typing import cast, Optional, TYPE_CHECKING, Union
 from typing_extensions import deprecated
 
 import torch
 import torch.distributed as dist
 from torch.distributed._state_dict_utils import STATE_DICT_TYPE
-from torch.distributed.checkpoint._async_executor import (  # noqa: TC001
-    _AsyncCheckpointExecutor,
-)
 from torch.distributed.checkpoint._async_process_executor import (
     _ProcessBasedAsyncCheckpointExecutor,
 )
@@ -38,6 +35,10 @@
 from .utils import _api_bc_check, _DistWrapper, _profile
 
 
+if TYPE_CHECKING:
+    from torch.distributed.checkpoint._async_executor import _AsyncCheckpointExecutor
+
+
 __all__ = [
     "save_state_dict",
     "save",

torch/fx/passes/_tensorify_python_scalars.py

@@ -2,7 +2,7 @@
 
 import logging
 import os
-from typing import Any, Union
+from typing import Any, TYPE_CHECKING, Union
 
 from sympy import Integer, Number, Symbol
 from sympy.logic.boolalg import BooleanAtom
@@ -13,16 +13,14 @@
 from torch._dynamo.symbolic_convert import TensorifyState
 from torch._dynamo.utils import get_metrics_context
 from torch._prims_common import get_computation_dtype
-from torch._subclasses import fake_tensor  # noqa: TCH001
 from torch._subclasses.fake_tensor import FakeTensor
 from torch._utils_internal import justknobs_check
 from torch.fx._utils import lazy_format_graph_code
-from torch.fx.experimental.symbolic_shapes import (  # noqa: TCH001
+from torch.fx.experimental.symbolic_shapes import (
     guard_scalar,
     has_free_symbols,
     ShapeEnv,
 )
-from torch.fx.graph_module import GraphModule  # noqa: TCH001
 
 # TODO: refactor
 from torch.fx.passes.runtime_assert import _get_sym_val
@@ -32,6 +30,11 @@
 from torch.utils._sympy.symbol import symbol_is_type, SymT
 
 
+if TYPE_CHECKING:
+    from torch._subclasses import fake_tensor
+    from torch.fx.graph_module import GraphModule
+
+
 __all__: list[str] = []
 
 log = logging.getLogger(__name__)

torch/onnx/_internal/exporter/_torchlib/ops/nn.py

@@ -1,7 +1,7 @@
 """torch.ops.aten operators under the `core` module."""
 # mypy: disable-error-code="misc,arg-type,type-arg,valid-type,assignment,return-value,type-var,operator,no-untyped-def,index"
 # pyrefly: ignore-errors
-# ruff: noqa: TCH001,TCH002
+# ruff: noqa: TC001,TC002
 # flake8: noqa: B950
 
 from __future__ import annotations

torch/onnx/_internal/fx/type_utils.py

@@ -14,7 +14,7 @@
 
 
 if TYPE_CHECKING:
-    import onnx.defs  # noqa: TCH004
+    import onnx.defs
 
 
 # Enable both TorchScriptTensor and torch.Tensor to be tested