manual fixes

Author: matsumotosan

src/lightning/fabric/connector.py

@@ -14,7 +14,7 @@
 import os
 from collections import Counter
 from collections.abc import Iterable
-from typing import Any, Union, cast, get_args
+from typing import Any, cast, get_args
 
 import torch
 
@@ -67,7 +67,7 @@
 from lightning.fabric.utilities.device_parser import _determine_root_gpu_device
 from lightning.fabric.utilities.imports import _IS_INTERACTIVE
 
-_PLUGIN_INPUT = Union[Precision, ClusterEnvironment, CheckpointIO]
+_PLUGIN_INPUT = Precision | ClusterEnvironment | CheckpointIO
 
 
 class _Connector:

src/lightning/fabric/plugins/precision/precision.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from contextlib import AbstractContextManager, nullcontext
-from typing import Any, Literal, Union
+from typing import Any, Literal
 
 from torch import Tensor
 from torch.nn import Module
@@ -33,7 +33,7 @@
     "32-true",
     "64-true",
 ]
-_PRECISION_INPUT = Union[_PRECISION_INPUT_INT, _PRECISION_INPUT_STR, _PRECISION_INPUT_STR_ALIAS]
+_PRECISION_INPUT = _PRECISION_INPUT_INT | _PRECISION_INPUT_STR | _PRECISION_INPUT_STR_ALIAS
 
 
 class Precision:

src/lightning/fabric/strategies/fsdp.py

@@ -73,8 +73,8 @@
     from torch.distributed.fsdp.wrap import ModuleWrapPolicy
     from torch.optim.lr_scheduler import _LRScheduler
 
-    _POLICY = Union[set[type[Module]], Callable[[Module, bool, int], bool], ModuleWrapPolicy]
-    _SHARDING_STRATEGY = Union[ShardingStrategy, Literal["FULL_SHARD", "SHARD_GRAD_OP", "NO_SHARD", "HYBRID_SHARD"]]
+    _POLICY = set[type[Module]] | Callable[[Module, bool, int], bool] | ModuleWrapPolicy
+    _SHARDING_STRATEGY = ShardingStrategy | Literal["FULL_SHARD", "SHARD_GRAD_OP", "NO_SHARD", "HYBRID_SHARD"]
 
 
 _FSDP_ALIASES = ("fsdp", "fsdp_cpu_offload")

src/lightning/fabric/strategies/xla_fsdp.py

@@ -16,7 +16,7 @@
 from contextlib import AbstractContextManager, ExitStack, nullcontext
 from functools import partial
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Literal, Optional
 
 import torch
 from torch import Tensor
@@ -49,7 +49,7 @@
     from torch_xla.distributed.parallel_loader import MpDeviceLoader
 
 _POLICY_SET = set[type[Module]]
-_POLICY = Union[_POLICY_SET, Callable[[Module, bool, int], bool]]
+_POLICY = _POLICY_SET | Callable[[Module, bool, int], bool]
 
 
 class XLAFSDPStrategy(ParallelStrategy, _Sharded):

src/lightning/fabric/utilities/types.py

@@ -16,11 +16,9 @@
 from pathlib import Path
 from typing import (
     Any,
-    Optional,
     Protocol,
     TypeAlias,
     TypeVar,
-    Union,
     runtime_checkable,
 )
 
@@ -30,9 +28,9 @@
 
 UntypedStorage: TypeAlias = torch.UntypedStorage
 
-_PATH = Union[str, Path]
-_DEVICE = Union[torch.device, str, int]
-_MAP_LOCATION_TYPE = Optional[_DEVICE | Callable[[UntypedStorage, str], UntypedStorage | None] | dict[_DEVICE, _DEVICE]]
+_PATH = str | Path
+_DEVICE = torch.device | str | int
+_MAP_LOCATION_TYPE = _DEVICE | Callable[[UntypedStorage, str], UntypedStorage | None] | dict[_DEVICE, _DEVICE] | None
 _PARAMETERS = Iterator[torch.nn.Parameter]
 
 if torch.distributed.is_available():

src/lightning/pytorch/cli.py

@@ -18,7 +18,7 @@
 from functools import partial, update_wrapper
 from pathlib import Path
 from types import MethodType
-from typing import Any, Optional, TypeVar, Union
+from typing import Any, TypeVar
 
 import torch
 import yaml
@@ -80,12 +80,12 @@ def __init__(self, optimizer: Optimizer, monitor: str, *args: Any, **kwargs: Any
 
 # LightningCLI requires the ReduceLROnPlateau defined here, thus it shouldn't accept the one from pytorch:
 LRSchedulerTypeTuple = (LRScheduler, ReduceLROnPlateau)
-LRSchedulerTypeUnion = Union[LRScheduler, ReduceLROnPlateau]
-LRSchedulerType = Union[type[LRScheduler], type[ReduceLROnPlateau]]
+LRSchedulerTypeUnion = LRScheduler | ReduceLROnPlateau
+LRSchedulerType = type[LRScheduler] | type[ReduceLROnPlateau]
 
 
 # Type aliases intended for convenience of CLI developers
-ArgsType = Optional[list[str] | dict[str, Any] | Namespace]
+ArgsType = list[str] | dict[str, Any] | Namespace | None
 OptimizerCallable = Callable[[Iterable], Optimizer]
 LRSchedulerCallable = Callable[[Optimizer], LRScheduler | ReduceLROnPlateau]
 
@@ -448,7 +448,7 @@ def add_default_arguments_to_parser(self, parser: LightningArgumentParser) -> No
         """Adds default arguments to the parser."""
         parser.add_argument(
             "--seed_everything",
-            type=Union[bool, int],
+            type=bool | int,
             default=self.seed_everything_default,
             help=(
                 "Set to an int to run seed_everything with this value before classes instantiation."

src/lightning/pytorch/core/module.py

@@ -27,7 +27,6 @@
     Any,
     Literal,
     Optional,
-    Union,
     cast,
     overload,
 )
@@ -89,9 +88,14 @@
 warning_cache = WarningCache()
 log = logging.getLogger(__name__)
 
-MODULE_OPTIMIZERS = Union[
-    Optimizer, LightningOptimizer, _FabricOptimizer, list[Optimizer], list[LightningOptimizer], list[_FabricOptimizer]
-]
+MODULE_OPTIMIZERS = (
+    Optimizer
+    | LightningOptimizer
+    | _FabricOptimizer
+    | list[Optimizer]
+    | list[LightningOptimizer]
+    | list[_FabricOptimizer]
+)
 
 
 class LightningModule(

src/lightning/pytorch/loops/training_epoch_loop.py

@@ -16,7 +16,7 @@
 import time
 from collections import OrderedDict
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any
 
 import torch
 from typing_extensions import override
@@ -38,7 +38,7 @@
 from lightning.pytorch.utilities.rank_zero import WarningCache, rank_zero_warn
 from lightning.pytorch.utilities.signature_utils import is_param_in_hook_signature
 
-_BATCH_OUTPUTS_TYPE = Optional[_OPTIMIZER_LOOP_OUTPUTS_TYPE | _MANUAL_LOOP_OUTPUTS_TYPE]
+_BATCH_OUTPUTS_TYPE = _OPTIMIZER_LOOP_OUTPUTS_TYPE | _MANUAL_LOOP_OUTPUTS_TYPE | None
 
 
 @dataclass

src/lightning/pytorch/plugins/__init__.py

@@ -1,5 +1,3 @@
-from typing import Union
-
 from lightning.fabric.plugins import CheckpointIO, ClusterEnvironment, TorchCheckpointIO, XLACheckpointIO
 from lightning.pytorch.plugins.io.async_plugin import AsyncCheckpointIO
 from lightning.pytorch.plugins.layer_sync import LayerSync, TorchSyncBatchNorm
@@ -13,7 +11,7 @@
 from lightning.pytorch.plugins.precision.transformer_engine import TransformerEnginePrecision
 from lightning.pytorch.plugins.precision.xla import XLAPrecision
 
-_PLUGIN_INPUT = Union[Precision, ClusterEnvironment, CheckpointIO, LayerSync]
+_PLUGIN_INPUT = Precision | ClusterEnvironment | CheckpointIO | LayerSync
 
 __all__ = [
     "AsyncCheckpointIO",

src/lightning/pytorch/profilers/pytorch.py

@@ -20,7 +20,7 @@
 from contextlib import AbstractContextManager
 from functools import lru_cache, partial
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any
 
 import torch
 from torch import Tensor, nn
@@ -42,7 +42,7 @@
 log = logging.getLogger(__name__)
 warning_cache = WarningCache()
 
-_PROFILER = Union[torch.profiler.profile, torch.autograd.profiler.profile, torch.autograd.profiler.emit_nvtx]
+_PROFILER = torch.profiler.profile | torch.autograd.profiler.profile | torch.autograd.profiler.emit_nvtx
 _KINETO_AVAILABLE = torch.profiler.kineto_available()