Fix .llm import trigger of .distributed depr warning

torch_geometric/distributed/__init__.py

@@ -1,5 +1,3 @@
-from warnings import warn
-
 from .dist_context import DistContext
 from .local_feature_store import LocalFeatureStore
 from .local_graph_store import LocalGraphStore
@@ -9,17 +7,6 @@
 from .dist_neighbor_loader import DistNeighborLoader
 from .dist_link_neighbor_loader import DistLinkNeighborLoader
 
-warn(
-    "`torch_geometric.distributed` has been deprecated since 2.7.0 and will "
-    "no longer be maintained. For distributed training, refer to our "
-    "tutorials on distributed training at "
-    "https://pytorch-geometric.readthedocs.io/en/latest/tutorial/distributed.html "  # noqa: E501
-    "or cuGraph examples at "
-    "https://github.com/rapidsai/cugraph-gnn/tree/main/python/cugraph-pyg/cugraph_pyg/examples",  # noqa: E501
-    stacklevel=2,
-    category=DeprecationWarning,
-)
-
 __all__ = classes = [
     'DistContext',
     'LocalFeatureStore',

torch_geometric/distributed/dist_link_neighbor_loader.py

@@ -1,4 +1,5 @@
 from typing import Callable, Dict, List, Optional, Tuple, Union
+from warnings import warn
 
 import torch
 
@@ -93,6 +94,15 @@ def __init__(
                 channel=channel,
                 concurrency=concurrency,
             )
+        else:
+            warn(
+                "`torch_geometric.distributed` has been deprecated since 2.7.0 and will "  # noqa: E501
+                "no longer be maintained. For distributed training, refer to our "  # noqa: E501
+                "tutorials on distributed training at "
+                "https://pytorch-geometric.readthedocs.io/en/latest/tutorial/distributed.html "  # noqa: E501
+                "or cuGraph examples at "
+                "https://github.com/rapidsai/cugraph-gnn/tree/main/python/cugraph-pyg/cugraph_pyg/examples",  # noqa: E501
+            )
 
         DistLoader.__init__(
             self,

torch_geometric/distributed/dist_neighbor_loader.py

@@ -1,4 +1,5 @@
 from typing import Callable, Dict, List, Optional, Tuple, Union
+from warnings import warn
 
 import torch
 
@@ -86,6 +87,15 @@ def __init__(
                 channel=channel,
                 concurrency=concurrency,
             )
+        else:
+            warn(
+                "`torch_geometric.distributed` has been deprecated since 2.7.0 and will "  # noqa: E501
+                "no longer be maintained. For distributed training, refer to our "  # noqa: E501
+                "tutorials on distributed training at "
+                "https://pytorch-geometric.readthedocs.io/en/latest/tutorial/distributed.html "  # noqa: E501
+                "or cuGraph examples at "
+                "https://github.com/rapidsai/cugraph-gnn/tree/main/python/cugraph-pyg/cugraph_pyg/examples",  # noqa: E501
+            )
 
         DistLoader.__init__(
             self,

torch_geometric/distributed/dist_neighbor_sampler.py

@@ -2,6 +2,7 @@
 import logging
 import math
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from warnings import warn
 
 import numpy as np
 import torch
@@ -81,6 +82,14 @@ def __init__(
         device: Optional[torch.device] = None,
         **kwargs,
     ):
+        warn(
+            "`torch_geometric.distributed` has been deprecated since 2.7.0 and will "  # noqa: E501
+            "no longer be maintained. For distributed training, refer to our "  # noqa: E501
+            "tutorials on distributed training at "
+            "https://pytorch-geometric.readthedocs.io/en/latest/tutorial/distributed.html "  # noqa: E501
+            "or cuGraph examples at "
+            "https://github.com/rapidsai/cugraph-gnn/tree/main/python/cugraph-pyg/cugraph_pyg/examples",  # noqa: E501
+        )
         self.current_ctx = current_ctx
 
         self.feature_store, self.graph_store = data

torch_geometric/llm/utils/feature_store.py

@@ -6,7 +6,7 @@
 from torch import Tensor
 
 from torch_geometric.data import Data, HeteroData
-from torch_geometric.distributed import LocalFeatureStore
+from torch_geometric.distributed.local_feature_store import LocalFeatureStore
 from torch_geometric.llm.utils.backend_utils import batch_knn
 from torch_geometric.sampler import HeteroSamplerOutput, SamplerOutput
 from torch_geometric.typing import InputNodes

torch_geometric/llm/utils/graph_store.py

@@ -4,7 +4,7 @@
 from torch import Tensor
 
 from torch_geometric.data import FeatureStore
-from torch_geometric.distributed import LocalGraphStore
+from torch_geometric.distributed.local_graph_store import LocalGraphStore
 from torch_geometric.sampler import (
     BidirectionalNeighborSampler,
     NodeSamplerInput,