Refactor imports in precision and strategy modules for consistency

Author: bhimrazy

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

@@ -12,13 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from contextlib import AbstractContextManager, nullcontext
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, get_args
 
 import torch
 from lightning_utilities.core.apply_func import apply_to_collection
 from torch import Tensor
 from torch.nn import Module
-from typing_extensions import get_args, override
+from typing_extensions import override
 
 from lightning.fabric.plugins.precision.precision import Precision
 from lightning.fabric.plugins.precision.utils import _convert_fp_tensor, _DtypeContextManager

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

@@ -12,10 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
-from typing import Any, Literal
+from typing import Any, Literal, get_args
 
 import torch
-from typing_extensions import get_args, override
+from typing_extensions import override
 
 from lightning.fabric.accelerators.xla import _XLA_AVAILABLE
 from lightning.fabric.plugins.precision.precision import Precision

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

tests/tests_pytorch/utilities/test_model_helpers.py

@@ -42,11 +42,10 @@ def test_is_overridden():
     reason="This test is ONLY relevant for the UNIFIED package",
 )
 def test_mixed_imports_unified():
-    from pytorch_lightning.callbacks import EarlyStopping as OldEarlyStopping
-    from pytorch_lightning.demos.boring_classes import BoringModel as OldBoringModel
-
     from lightning.pytorch.utilities.compile import _maybe_unwrap_optimized as new_unwrap
     from lightning.pytorch.utilities.model_helpers import is_overridden as new_is_overridden
+    from pytorch_lightning.callbacks import EarlyStopping as OldEarlyStopping
+    from pytorch_lightning.demos.boring_classes import BoringModel as OldBoringModel
 
     model = OldBoringModel()
     with pytest.raises(TypeError, match=r"`pytorch_lightning` object \(BoringModel\) to a `lightning.pytorch`"):