ValueError("Expected a parent") when using new Lightning imports and PyG Lightning DataModule wrappers

[nathanpainchaud]

🐛 Describe the bug

Bug

When using modern Lightning imports (i.e. from lightning import ...) rather than old imports (i.e. from pytorch_lightning import ...) in a project mixing Lightning and PyG, the code crashes on a non-intuitive ValueError("Expected a parent"). I suspect the exact line that raises this message depends on the user application.

This seems to be caused by Lightning sometimes (but not always!) failing when import styles are mixed together (see this issue for context). Given that PyG's Lightning wrappers use the old style, but most other codebases use the modern style, a conflict happens.

Below is the smallest working example I could come up with that reproduces the issue:

import torch
from lightning import LightningModule, Trainer
from torch.nn import functional as F
from torch_geometric.data.lightning import LightningNodeData
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCN


class GCNLitModule(LightningModule):
    def __init__(self, in_features: int, num_classes: int):
        super().__init__()
        self.model = GCN(in_channels=in_features, hidden_channels=16, num_layers=2, out_channels=num_classes)

    def forward(self, batch):
        return self.model(batch.x, batch.edge_index)

    def training_step(self, batch):
        out = self.forward(batch)
        loss = F.nll_loss(out, batch.y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.model.parameters(), lr=0.01)


if __name__ == "__main__":
    pyg_dataset = Planetoid(root="../data", name="Cora")
    lit_model = GCNLitModule(in_features=pyg_dataset.num_features, num_classes=pyg_dataset.num_classes)
    lit_datamodule = LightningNodeData(pyg_dataset[0], batch_size=32, loader="full")
    trainer = Trainer(max_epochs=10)
    trainer.fit(model=lit_model, datamodule=lit_datamodule)

Alternative considered

In this specific example, the error can be fixed by reverting back to the old import style:

- from lightning import LightningModule, Trainer
+ from pytorch_lightning import LightningModule, Trainer

However, in more complex projects, e.g. with other dependencies using Lightning, it's not realistic (and often outside of the user's power) to revert all imports to the old style. For example, in the real project where I first encountered this error, troubleshooting lead me to edit imports in dependencies...

Suggested fix

Given that the modern Lightning import style is not so new anymore (and rather widely adopted from I could see), I would suggest PyG update how they import Lightning. I was able to fix the example above (and the issue in my real project) by simply changing PyG's Lightning imports like below:

try:
-    from pytorch_lightning import LightningDataModule as PLLightningDataModule
+    from lightning import LightningDataModule as PLLightningDataModule
    no_pytorch_lightning = False
except ImportError:
    PLLightningDataModule = object  # type: ignore
    no_pytorch_lightning = True

Versions

I was not able to obtain the full info about my environment because I'm using uv and the script crashes on a uv environment. But I am experiencing this with Python 3.12, and here is the relevant info about the torch-related dependencies:

lightning==2.4.0
lightning-utilities==0.11.9
torch==2.5.1
torch-geometric==2.6.1
pytorch-lightning==2.4.0