Check for mixed new and old style imports

src/lightning/pytorch/utilities/compile.py

@@ -18,6 +18,7 @@
 import lightning.pytorch as pl
 from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0, _TORCH_GREATER_EQUAL_2_1
 from lightning.pytorch.strategies import DDPStrategy, FSDPStrategy, SingleDeviceStrategy, Strategy
+from lightning.pytorch.utilities.model_helpers import _check_mixed_imports
 
 
 def from_compiled(model: "torch._dynamo.OptimizedModule") -> "pl.LightningModule":
@@ -42,6 +43,7 @@ def from_compiled(model: "torch._dynamo.OptimizedModule") -> "pl.LightningModule
     orig_module = model._orig_mod
 
     if not isinstance(orig_module, pl.LightningModule):
+        _check_mixed_imports(model)
         raise ValueError(
             f"`model` is expected to be a compiled LightingModule. Found a `{type(orig_module).__name__}` instead"
         )
@@ -114,6 +116,7 @@ def to_uncompiled(model: Union["pl.LightningModule", "torch._dynamo.OptimizedMod
 def _maybe_unwrap_optimized(model: object) -> "pl.LightningModule":
     if not _TORCH_GREATER_EQUAL_2_0:
         if not isinstance(model, pl.LightningModule):
+            _check_mixed_imports(model)
             raise TypeError(f"`model` must be a `LightningModule`, got `{type(model).__qualname__}`")
         return model
     from torch._dynamo import OptimizedModule
@@ -122,6 +125,7 @@ def _maybe_unwrap_optimized(model: object) -> "pl.LightningModule":
         return from_compiled(model)
     if isinstance(model, pl.LightningModule):
         return model
+    _check_mixed_imports(model)
     raise TypeError(
         f"`model` must be a `LightningModule` or `torch._dynamo.OptimizedModule`, got `{type(model).__qualname__}`"
     )

src/lightning/pytorch/utilities/model_helpers.py

@@ -31,6 +31,7 @@ def is_overridden(method_name: str, instance: Optional[object] = None, parent: O
         elif isinstance(instance, pl.Callback):
             parent = pl.Callback
         if parent is None:
+            _check_mixed_imports(instance)
             raise ValueError("Expected a parent")
 
     from lightning_utilities.core.overrides import is_overridden as _is_overridden
@@ -51,3 +52,19 @@ def get_torchvision_model(model_name: str, **kwargs: Any) -> nn.Module:
     if torchvision_greater_equal_0_14:
         return models.get_model(model_name, **kwargs)
     return getattr(models, model_name)(**kwargs)
+
+
+def _check_mixed_imports(instance: object) -> None:
+    old, new = "pytorch_" + "lightning", "lightning." + "pytorch"
+    klass = type(instance)
+    module = klass.__module__
+    if module.startswith(old) and __name__.startswith(new):
+        pass
+    elif module.startswith(new) and __name__.startswith(old):
+        old, new = new, old
+    else:
+        return
+    raise TypeError(
+        f"You passed a `{old}` object ({type(instance).__qualname__}) to a `{new}`"
+        " Trainer. Please switch to a single import style."
+    )

tests/tests_pytorch/graveyard/test_legacy_import_unpickler.py

@@ -20,9 +20,7 @@ def _list_sys_modules(pattern: str) -> str:
 
 
 @pytest.mark.parametrize("pl_version", LEGACY_BACK_COMPATIBLE_PL_VERSIONS)
-@pytest.mark.skipif(
-    not module_available("lightning_pytorch"), reason="This test is ONLY relevant for the STANDALONE package"
-)
+@pytest.mark.skipif(module_available("lightning"), reason="This test is ONLY relevant for the STANDALONE package")
 def test_imports_standalone(pl_version: str):
     assert any(
         key.startswith("pytorch_lightning") for key in sys.modules

tests/tests_pytorch/utilities/test_model_helpers.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import pytest
+from lightning_utilities import module_available
 
 from lightning.pytorch import LightningDataModule
 from lightning.pytorch.demos.boring_classes import BoringDataModule, BoringModel
@@ -30,3 +31,21 @@ def test_is_overridden():
     assert is_overridden("training_step", model)
     datamodule = BoringDataModule()
     assert is_overridden("train_dataloader", datamodule)
+
+
+@pytest.mark.skipif(
+    not module_available("lightning") or not module_available("pytorch_lightning"),
+    reason="This test is ONLY relevant for the UNIFIED package",
+)
+def test_mixed_imports_unified():
+    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`"):
+        new_unwrap(model)
+
+    with pytest.raises(TypeError, match=r"`pytorch_lightning` object \(EarlyStopping\) to a `lightning.pytorch`"):
+        new_is_overridden("on_fit_start", OldEarlyStopping("foo"))