pyg-nightly 2.7.0.dev20241009__py3-none-any.whl → 2.8.0.dev20251228__py3-none-any.whl

torch_geometric/testing/decorators.py

@@ -10,10 +10,16 @@ from packaging.requirements import Requirement
 from packaging.version import Version
 
 import torch_geometric
+import torch_geometric.typing
 from torch_geometric.typing import WITH_METIS, WITH_PYG_LIB, WITH_TORCH_SPARSE
 from torch_geometric.visualization.graph import has_graphviz
 
 
+def is_rag_test() -> bool:
+    r"""Whether to run the RAG test suite."""
+    return os.getenv('RAG_TEST', '0') == '1'
+
+
 def is_full_test() -> bool:
     r"""Whether to run the full but time-consuming test suite."""
     return os.getenv('FULL_TEST', '0') == '1'
@@ -32,8 +38,8 @@ def onlyFullTest(func: Callable) -> Callable:
 
 def is_distributed_test() -> bool:
     r"""Whether to run the distributed test suite."""
-    return ((is_full_test() or os.getenv('DIST_TEST', '0') == '1')
-            and sys.platform == 'linux' and has_package('pyg_lib'))
+    return (os.getenv('DIST_TEST', '0') == '1' and sys.platform == 'linux'
+            and has_package('pyg_lib'))
 
 
 def onlyDistributedTest(func: Callable) -> Callable:
@@ -203,6 +209,18 @@ def withPackage(*args: str) -> Callable:
     return decorator
 
 
+def onlyRAG(func: Callable) -> Callable:
+    r"""A decorator to specify that this function belongs to the RAG test
+    suite.
+    """
+    import pytest
+    func = pytest.mark.rag(func)
+    return pytest.mark.skipif(
+        not is_rag_test(),
+        reason="RAG tests are disabled",
+    )(func)
+
+
 def withCUDA(func: Callable) -> Callable:
     r"""A decorator to test both on CPU and CUDA (if available)."""
     import pytest
@@ -234,8 +252,9 @@ def withDevice(func: Callable) -> Callable:
     if device:
         backend = os.getenv('TORCH_BACKEND')
         if backend is None:
-            warnings.warn(f"Please specify the backend via 'TORCH_BACKEND' in"
-                          f"order to test against '{device}'")
+            warnings.warn(
+                f"Please specify the backend via 'TORCH_BACKEND' in"
+                f"order to test against '{device}'", stacklevel=2)
         else:
             import_module(backend)
             devices.append(pytest.param(torch.device(device), id=device))
@@ -250,7 +269,7 @@ def withMETIS(func: Callable) -> Callable:
     with_metis = WITH_METIS
 
     if with_metis:
-        try:  # Test that METIS can succesfully execute:
+        try:  # Test that METIS can successfully execute:
             # TODO Using `pyg-lib` metis partitioning leads to some weird bugs
             # in the # CI. As such, we require `torch-sparse` for now.
             rowptr = torch.tensor([0, 2, 4, 6])
@@ -265,6 +284,17 @@ def withMETIS(func: Callable) -> Callable:
     )(func)
 
 
+def withHashTensor(func: Callable) -> Callable:
+    r"""A decorator to only test in case :class:`HashTensor` is available."""
+    import pytest
+
+    return pytest.mark.skipif(
+        not torch_geometric.typing.WITH_CPU_HASH_MAP
+        and not has_package('pandas'),
+        reason="HashTensor dependencies not available",
+    )(func)
+
+
 def disableExtensions(func: Callable) -> Callable:
     r"""A decorator to temporarily disable the usage of the
     :obj:`torch_scatter`, :obj:`torch_sparse` and :obj:`pyg_lib` extension

torch_geometric/testing/distributed.py

@@ -73,7 +73,7 @@ def assert_run_mproc(
     ]
     results = []
 
-    for p, q in zip(procs, queues):
+    for p, _ in zip(procs, queues):
         p.start()
 
     for p, q in zip(procs, queues):

torch_geometric/transforms/__init__.py

@@ -37,6 +37,7 @@ from .rooted_subgraph import RootedEgoNets, RootedRWSubgraph
 from .largest_connected_components import LargestConnectedComponents
 from .virtual_node import VirtualNode
 from .add_positional_encoding import AddLaplacianEigenvectorPE, AddRandomWalkPE
+from .add_gpse import AddGPSE
 from .feature_propagation import FeaturePropagation
 from .half_hop import HalfHop
 
@@ -108,6 +109,7 @@ graph_transforms = [
     'VirtualNode',
     'AddLaplacianEigenvectorPE',
     'AddRandomWalkPE',
+    'AddGPSE',
     'FeaturePropagation',
     'HalfHop',
 ]

torch_geometric/transforms/add_gpse.py

@@ -0,0 +1,49 @@
+from typing import Any
+
+from torch.nn import Module
+
+from torch_geometric.data import Data
+from torch_geometric.data.datapipes import functional_transform
+from torch_geometric.transforms import BaseTransform, VirtualNode
+
+
+@functional_transform('add_gpse')
+class AddGPSE(BaseTransform):
+    r"""Adds the GPSE encoding from the `"Graph Positional and Structural
+    Encoder" <https://arxiv.org/abs/2307.07107>`_ paper to the given graph
+    (functional name: :obj:`add_gpse`).
+    To be used with a :class:`~torch_geometric.nn.GPSE` model, which generates
+    the actual encodings.
+
+    Args:
+        model (Module): The pre-trained GPSE model.
+        use_vn (bool, optional): Whether to use virtual nodes.
+            (default: :obj:`True`)
+        rand_type (str, optional): Type of random features to use. Options are
+            :obj:`NormalSE`, :obj:`UniformSE`, :obj:`BernoulliSE`.
+            (default: :obj:`NormalSE`)
+
+    """
+    def __init__(
+        self,
+        model: Module,
+        use_vn: bool = True,
+        rand_type: str = 'NormalSE',
+    ):
+        self.model = model
+        self.use_vn = use_vn
+        self.vn = VirtualNode()
+        self.rand_type = rand_type
+
+    def forward(self, data: Data) -> Any:
+        pass
+
+    def __call__(self, data: Data) -> Data:
+        from torch_geometric.nn.models.gpse import gpse_process
+
+        data_vn = self.vn(data.clone()) if self.use_vn else data.clone()
+        batch_out = gpse_process(self.model, data_vn, 'NormalSE', self.use_vn)
+        batch_out = batch_out.to('cpu', non_blocking=True)
+        data.pestat_GPSE = batch_out[:-1] if self.use_vn else batch_out
+
+        return data

torch_geometric/transforms/add_metapaths.py

@@ -108,13 +108,15 @@ class AddMetaPaths(BaseTransform):
         **kwargs: bool,
     ) -> None:
         if 'drop_orig_edges' in kwargs:
-            warnings.warn("'drop_orig_edges' is deprecated. Use "
-                          "'drop_orig_edge_types' instead")
+            warnings.warn(
+                "'drop_orig_edges' is deprecated. Use "
+                "'drop_orig_edge_types' instead", stacklevel=2)
             drop_orig_edge_types = kwargs['drop_orig_edges']
 
         if 'drop_unconnected_nodes' in kwargs:
-            warnings.warn("'drop_unconnected_nodes' is deprecated. Use "
-                          "'drop_unconnected_node_types' instead")
+            warnings.warn(
+                "'drop_unconnected_nodes' is deprecated. Use "
+                "'drop_unconnected_node_types' instead", stacklevel=2)
             drop_unconnected_node_types = kwargs['drop_unconnected_nodes']
 
         for path in metapaths:
@@ -144,7 +146,7 @@ class AddMetaPaths(BaseTransform):
             if self.max_sample is not None:
                 edge_index, edge_weight = self._sample(edge_index, edge_weight)
 
-            for i, edge_type in enumerate(metapath[1:]):
+            for edge_type in metapath[1:]:
                 edge_index2, edge_weight2 = self._edge_index(data, edge_type)
 
                 edge_index, edge_weight = edge_index.matmul(
@@ -276,7 +278,7 @@ class AddRandomMetaPaths(BaseTransform):
             row = start = torch.randperm(num_nodes)[:num_starts].repeat(
                 self.walks_per_node[j])
 
-            for i, edge_type in enumerate(metapath):
+            for edge_type in metapath:
                 edge_index = EdgeIndex(
                     data[edge_type].edge_index,
                     sparse_size=data[edge_type].size(),

torch_geometric/transforms/add_positional_encoding.py

@@ -92,12 +92,12 @@ class AddLaplacianEigenvectorPE(BaseTransform):
             from numpy.linalg import eig, eigh
             eig_fn = eig if not self.is_undirected else eigh
 
-            eig_vals, eig_vecs = eig_fn(L.todense())  # type: ignore
+            eig_vals, eig_vecs = eig_fn(L.todense())
         else:
             from scipy.sparse.linalg import eigs, eigsh
             eig_fn = eigs if not self.is_undirected else eigsh
 
-            eig_vals, eig_vecs = eig_fn(  # type: ignore
+            eig_vals, eig_vecs = eig_fn(
                 L,
                 k=self.k + 1,
                 which='SR' if not self.is_undirected else 'SA',

torch_geometric/transforms/base_transform.py

@@ -1,5 +1,5 @@
 import copy
-from abc import ABC
+from abc import ABC, abstractmethod
 from typing import Any
 
 
@@ -31,6 +31,7 @@ class BaseTransform(ABC):
         # Shallow-copy the data so that we prevent in-place data modification.
         return self.forward(copy.copy(data))
 
+    @abstractmethod
     def forward(self, data: Any) -> Any:
         pass
 

torch_geometric/transforms/delaunay.py

@@ -1,3 +1,5 @@
+from typing import List
+
 import torch
 
 from torch_geometric.data import Data
@@ -5,30 +7,78 @@ from torch_geometric.data.datapipes import functional_transform
 from torch_geometric.transforms import BaseTransform
 
 
+class _QhullTransform(BaseTransform):
+    r"""Q-hull implementation of delaunay triangulation."""
+    def forward(self, data: Data) -> Data:
+        assert data.pos is not None
+        import scipy.spatial
+
+        pos = data.pos.cpu().numpy()
+        tri = scipy.spatial.Delaunay(pos, qhull_options='QJ')
+        face = torch.from_numpy(tri.simplices)
+
+        data.face = face.t().contiguous().to(data.pos.device, torch.long)
+        return data
+
+
+class _ShullTransform(BaseTransform):
+    r"""Sweep-hull implementation of delaunay triangulation."""
+    def forward(self, data: Data) -> Data:
+        assert data.pos is not None
+        from torch_delaunay.functional import shull2d
+
+        face = shull2d(data.pos.cpu())
+        data.face = face.t().contiguous().to(data.pos.device)
+        return data
+
+
+class _SequentialTransform(BaseTransform):
+    r"""Runs the first successful transformation.
+
+    All intermediate exceptions are suppressed except the last.
+    """
+    def __init__(self, transforms: List[BaseTransform]) -> None:
+        assert len(transforms) > 0
+        self.transforms = transforms
+
+    def forward(self, data: Data) -> Data:
+        for i, transform in enumerate(self.transforms):
+            try:
+                return transform.forward(data)
+            except ImportError as e:
+                if i == len(self.transforms) - 1:
+                    raise e
+        return data
+
+
 @functional_transform('delaunay')
 class Delaunay(BaseTransform):
     r"""Computes the delaunay triangulation of a set of points
     (functional name: :obj:`delaunay`).
+
+    .. hint::
+        Consider installing the
+        `torch_delaunay <https://github.com/ybubnov/torch_delaunay>`_ package
+        to speed up computation.
     """
-    def forward(self, data: Data) -> Data:
-        import scipy.spatial
+    def __init__(self) -> None:
+        self._transform = _SequentialTransform([
+            _ShullTransform(),
+            _QhullTransform(),
+        ])
 
+    def forward(self, data: Data) -> Data:
         assert data.pos is not None
+        device = data.pos.device
 
         if data.pos.size(0) < 2:
-            data.edge_index = torch.tensor([], dtype=torch.long,
-                                           device=data.pos.device).view(2, 0)
-        if data.pos.size(0) == 2:
-            data.edge_index = torch.tensor([[0, 1], [1, 0]], dtype=torch.long,
-                                           device=data.pos.device)
+            data.edge_index = torch.empty(2, 0, dtype=torch.long,
+                                          device=device)
+        elif data.pos.size(0) == 2:
+            data.edge_index = torch.tensor([[0, 1], [1, 0]], device=device)
         elif data.pos.size(0) == 3:
-            data.face = torch.tensor([[0], [1], [2]], dtype=torch.long,
-                                     device=data.pos.device)
-        if data.pos.size(0) > 3:
-            pos = data.pos.cpu().numpy()
-            tri = scipy.spatial.Delaunay(pos, qhull_options='QJ')
-            face = torch.from_numpy(tri.simplices)
-
-            data.face = face.t().contiguous().to(data.pos.device, torch.long)
+            data.face = torch.tensor([[0], [1], [2]], device=device)
+        else:
+            data = self._transform.forward(data)
 
         return data

torch_geometric/transforms/face_to_edge.py

@@ -8,8 +8,15 @@ from torch_geometric.utils import to_undirected
 
 @functional_transform('face_to_edge')
 class FaceToEdge(BaseTransform):
-    r"""Converts mesh faces :obj:`[3, num_faces]` to edge indices
-    :obj:`[2, num_edges]` (functional name: :obj:`face_to_edge`).
+    r"""Converts mesh faces of shape :obj:`[3, num_faces]` or
+    :obj:`[4, num_faces]` to edge indices of shape :obj:`[2, num_edges]`
+    (functional name: :obj:`face_to_edge`).
+
+    This transform supports both 2D triangular faces, represented by a
+    tensor of shape :obj:`[3, num_faces]`, and 3D tetrahedral mesh faces,
+    represented by a tensor of shape :obj:`[4, num_faces]`. It will convert
+    these faces into edge indices, where each edge is defined by the indices
+    of its two endpoints.
 
     Args:
         remove_faces (bool, optional): If set to :obj:`False`, the face tensor
@@ -22,7 +29,29 @@ class FaceToEdge(BaseTransform):
         if hasattr(data, 'face'):
             assert data.face is not None
             face = data.face
-            edge_index = torch.cat([face[:2], face[1:], face[::2]], dim=1)
+
+            if face.size(0) not in [3, 4]:
+                raise RuntimeError(f"Expected 'face' tensor with shape "
+                                   f"[3, num_faces] or [4, num_faces] "
+                                   f"(got {list(face.size())})")
+
+            if face.size()[0] == 3:
+                edge_index = torch.cat([
+                    face[:2],
+                    face[1:],
+                    face[::2],
+                ], dim=1)
+            else:
+                assert face.size()[0] == 4
+                edge_index = torch.cat([
+                    face[:2],
+                    face[1:3],
+                    face[2:4],
+                    face[::2],
+                    face[1::2],
+                    face[::3],
+                ], dim=1)
+
             edge_index = to_undirected(edge_index, num_nodes=data.num_nodes)
 
             data.edge_index = edge_index

torch_geometric/transforms/gdc.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 import numpy as np
 import torch
@@ -78,18 +78,17 @@ class GDC(BaseTransform):
         self_loop_weight: float = 1.,
         normalization_in: str = 'sym',
         normalization_out: str = 'col',
-        diffusion_kwargs: Dict[str, Any] = dict(method='ppr', alpha=0.15),
-        sparsification_kwargs: Dict[str, Any] = dict(
-            method='threshold',
-            avg_degree=64,
-        ),
+        diffusion_kwargs: Optional[Dict[str, Any]] = None,
+        sparsification_kwargs: Optional[Dict[str, Any]] = None,
         exact: bool = True,
     ) -> None:
         self.self_loop_weight = self_loop_weight
         self.normalization_in = normalization_in
         self.normalization_out = normalization_out
-        self.diffusion_kwargs = diffusion_kwargs
-        self.sparsification_kwargs = sparsification_kwargs
+        self.diffusion_kwargs = diffusion_kwargs or dict(
+            method='ppr', alpha=0.15)
+        self.sparsification_kwargs = sparsification_kwargs or dict(
+            method='threshold', avg_degree=64)
         self.exact = exact
 
         if self_loop_weight:

torch_geometric/transforms/largest_connected_components.py

@@ -47,7 +47,7 @@ class LargestConnectedComponents(BaseTransform):
             return data
 
         _, count = np.unique(component, return_counts=True)
-        subset_np = np.in1d(component, count.argsort()[-self.num_components:])
+        subset_np = np.isin(component, count.argsort()[-self.num_components:])
         subset = torch.from_numpy(subset_np)
         subset = subset.to(data.edge_index.device, torch.bool)
 

torch_geometric/transforms/mask.py

@@ -19,7 +19,11 @@ def get_attrs_with_suffix(
     return [key for key in store.keys() if key.endswith(suffix)]
 
 
-def get_mask_size(attr: str, store: BaseStorage, size: Optional[int]) -> int:
+def get_mask_size(
+    attr: str,
+    store: BaseStorage,
+    size: Optional[int],
+) -> Optional[int]:
     if size is not None:
         return size
     return store.num_edges if store.is_edge_attr(attr) else store.num_nodes

torch_geometric/transforms/normalize_features.py

@@ -1,4 +1,4 @@
-from typing import List, Union
+from typing import List, Optional, Union
 
 from torch_geometric.data import Data, HeteroData
 from torch_geometric.data.datapipes import functional_transform
@@ -14,8 +14,8 @@ class NormalizeFeatures(BaseTransform):
         attrs (List[str]): The names of attributes to normalize.
             (default: :obj:`["x"]`)
     """
-    def __init__(self, attrs: List[str] = ["x"]):
-        self.attrs = attrs
+    def __init__(self, attrs: Optional[List[str]] = None) -> None:
+        self.attrs = attrs or ["x"]
 
     def forward(
         self,

torch_geometric/transforms/random_link_split.py

@@ -245,7 +245,7 @@ class RandomLinkSplit(BaseTransform):
                 warnings.warn(
                     f"There are not enough negative edges to satisfy "
                     "the provided sampling ratio. The ratio will be "
-                    f"adjusted to {ratio:.2f}.")
+                    f"adjusted to {ratio:.2f}.", stacklevel=2)
                 num_neg_train = int((num_neg_train / num_neg) * num_neg_found)
                 num_neg_val = int((num_neg_val / num_neg) * num_neg_found)
                 num_neg_test = num_neg_found - num_neg_train - num_neg_val

torch_geometric/transforms/remove_duplicated_edges.py

@@ -1,4 +1,4 @@
-from typing import List, Union
+from typing import List, Optional, Union
 
 from torch_geometric.data import Data, HeteroData
 from torch_geometric.data.datapipes import functional_transform
@@ -22,9 +22,11 @@ class RemoveDuplicatedEdges(BaseTransform):
     """
     def __init__(
         self,
-        key: Union[str, List[str]] = ['edge_attr', 'edge_weight'],
+        key: Optional[Union[str, List[str]]] = None,
         reduce: str = "add",
     ) -> None:
+        key = key or ['edge_attr', 'edge_weight']
+
         if isinstance(key, str):
             key = [key]
 

torch_geometric/transforms/rooted_subgraph.py

@@ -94,7 +94,7 @@ class RootedSubgraph(BaseTransform, ABC):
         arange = torch.arange(n_id.size(0), device=data.edge_index.device)
         node_map = data.edge_index.new_ones(num_nodes, num_nodes)
         node_map[n_sub_batch, n_id] = arange
-        sub_edge_index += (arange * data.num_nodes)[e_sub_batch]
+        sub_edge_index += (arange * num_nodes)[e_sub_batch]
         sub_edge_index = node_map.view(-1)[sub_edge_index]
 
         return sub_edge_index, n_id, e_id, n_sub_batch, e_sub_batch

torch_geometric/typing.py

@@ -3,7 +3,7 @@ import os
 import sys
 import typing
 import warnings
-from typing import Any, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Dict, List, Optional, Set, Tuple, TypeAlias, Union
 
 import numpy as np
 import torch
@@ -15,8 +15,9 @@ WITH_PT22 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 2
 WITH_PT23 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 3
 WITH_PT24 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 4
 WITH_PT25 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 5
-WITH_PT111 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 11
-WITH_PT112 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 12
+WITH_PT26 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 6
+WITH_PT27 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 7
+WITH_PT28 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 8
 WITH_PT113 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 13
 
 WITH_WINDOWS = os.name == 'nt'
@@ -63,10 +64,21 @@ try:
         pyg_lib.sampler.neighbor_sample).parameters)
     WITH_WEIGHTED_NEIGHBOR_SAMPLE = ('edge_weight' in inspect.signature(
         pyg_lib.sampler.neighbor_sample).parameters)
+    try:
+        torch.classes.pyg.CPUHashMap  # noqa: B018
+        WITH_CPU_HASH_MAP = True
+    except Exception:
+        WITH_CPU_HASH_MAP = False
+    try:
+        torch.classes.pyg.CUDAHashMap  # noqa: B018
+        WITH_CUDA_HASH_MAP = True
+    except Exception:
+        WITH_CUDA_HASH_MAP = False
 except Exception as e:
     if not isinstance(e, ImportError):  # pragma: no cover
-        warnings.warn(f"An issue occurred while importing 'pyg-lib'. "
-                      f"Disabling its usage. Stacktrace: {e}")
+        warnings.warn(
+            f"An issue occurred while importing 'pyg-lib'. "
+            f"Disabling its usage. Stacktrace: {e}", stacklevel=2)
     pyg_lib = object
     WITH_PYG_LIB = False
     WITH_GMM = False
@@ -77,14 +89,41 @@ except Exception as e:
     WITH_METIS = False
     WITH_EDGE_TIME_NEIGHBOR_SAMPLE = False
     WITH_WEIGHTED_NEIGHBOR_SAMPLE = False
+    WITH_CPU_HASH_MAP = False
+    WITH_CUDA_HASH_MAP = False
+
+if WITH_CPU_HASH_MAP:
+    CPUHashMap: TypeAlias = torch.classes.pyg.CPUHashMap  # type: ignore[name-defined]  # noqa: E501
+else:
+
+    class CPUHashMap:  # type: ignore
+        def __init__(self, key: Tensor) -> None:
+            raise ImportError("'CPUHashMap' requires 'pyg-lib'")
+
+        def get(self, query: Tensor) -> Tensor:
+            raise ImportError("'CPUHashMap' requires 'pyg-lib'")
+
+
+if WITH_CUDA_HASH_MAP:
+    CUDAHashMap: TypeAlias = torch.classes.pyg.CUDAHashMap  # type: ignore[name-defined]  # noqa: E501
+else:
+
+    class CUDAHashMap:  # type: ignore
+        def __init__(self, key: Tensor) -> None:
+            raise ImportError("'CUDAHashMap' requires 'pyg-lib'")
+
+        def get(self, query: Tensor) -> Tensor:
+            raise ImportError("'CUDAHashMap' requires 'pyg-lib'")
+
 
 try:
     import torch_scatter  # noqa
     WITH_TORCH_SCATTER = True
 except Exception as e:
     if not isinstance(e, ImportError):  # pragma: no cover
-        warnings.warn(f"An issue occurred while importing 'torch-scatter'. "
-                      f"Disabling its usage. Stacktrace: {e}")
+        warnings.warn(
+            f"An issue occurred while importing 'torch-scatter'. "
+            f"Disabling its usage. Stacktrace: {e}", stacklevel=2)
     torch_scatter = object
     WITH_TORCH_SCATTER = False
 
@@ -94,8 +133,9 @@ try:
     WITH_TORCH_CLUSTER_BATCH_SIZE = 'batch_size' in torch_cluster.knn.__doc__
 except Exception as e:
     if not isinstance(e, ImportError):  # pragma: no cover
-        warnings.warn(f"An issue occurred while importing 'torch-cluster'. "
-                      f"Disabling its usage. Stacktrace: {e}")
+        warnings.warn(
+            f"An issue occurred while importing 'torch-cluster'. "
+            f"Disabling its usage. Stacktrace: {e}", stacklevel=2)
     WITH_TORCH_CLUSTER = False
     WITH_TORCH_CLUSTER_BATCH_SIZE = False
 
@@ -112,7 +152,7 @@ except Exception as e:
     if not isinstance(e, ImportError):  # pragma: no cover
         warnings.warn(
             f"An issue occurred while importing 'torch-spline-conv'. "
-            f"Disabling its usage. Stacktrace: {e}")
+            f"Disabling its usage. Stacktrace: {e}", stacklevel=2)
     WITH_TORCH_SPLINE_CONV = False
 
 try:
@@ -121,8 +161,9 @@ try:
     WITH_TORCH_SPARSE = True
 except Exception as e:
     if not isinstance(e, ImportError):  # pragma: no cover
-        warnings.warn(f"An issue occurred while importing 'torch-sparse'. "
-                      f"Disabling its usage. Stacktrace: {e}")
+        warnings.warn(
+            f"An issue occurred while importing 'torch-sparse'. "
+            f"Disabling its usage. Stacktrace: {e}", stacklevel=2)
     WITH_TORCH_SPARSE = False
 
     class SparseStorage:  # type: ignore
@@ -306,6 +347,8 @@ class EdgeTypeStr(str):
     r"""A helper class to construct serializable edge types by merging an edge
     type tuple into a single string.
     """
+    edge_type: tuple[str, str, str]
+
     def __new__(cls, *args: Any) -> 'EdgeTypeStr':
         if isinstance(args[0], (list, tuple)):
             # Unwrap `EdgeType((src, rel, dst))` and `EdgeTypeStr((src, dst))`:
@@ -313,27 +356,37 @@ class EdgeTypeStr(str):
 
         if len(args) == 1 and isinstance(args[0], str):
             arg = args[0]  # An edge type string was passed.
+            edge_type = tuple(arg.split(EDGE_TYPE_STR_SPLIT))
+            if len(edge_type) != 3:
+                raise ValueError(f"Cannot convert the edge type '{arg}' to a "
+                                 f"tuple since it holds invalid characters")
 
         elif len(args) == 2 and all(isinstance(arg, str) for arg in args):
             # A `(src, dst)` edge type was passed - add `DEFAULT_REL`:
-            arg = EDGE_TYPE_STR_SPLIT.join((args[0], DEFAULT_REL, args[1]))
+            edge_type = (args[0], DEFAULT_REL, args[1])
+            arg = EDGE_TYPE_STR_SPLIT.join(edge_type)
 
         elif len(args) == 3 and all(isinstance(arg, str) for arg in args):
             # A `(src, rel, dst)` edge type was passed:
+            edge_type = tuple(args)
             arg = EDGE_TYPE_STR_SPLIT.join(args)
 
         else:
             raise ValueError(f"Encountered invalid edge type '{args}'")
 
-        return str.__new__(cls, arg)
+        out = str.__new__(cls, arg)
+        out.edge_type = edge_type  # type: ignore
+        return out
 
     def to_tuple(self) -> EdgeType:
         r"""Returns the original edge type."""
-        out = tuple(self.split(EDGE_TYPE_STR_SPLIT))
-        if len(out) != 3:
+        if len(self.edge_type) != 3:
             raise ValueError(f"Cannot convert the edge type '{self}' to a "
                              f"tuple since it holds invalid characters")
-        return out
+        return self.edge_type
+
+    def __reduce__(self) -> tuple[Any, Any]:
+        return (self.__class__, (self.edge_type, ))
 
 
 # There exist some short-cuts to query edge-types (given that the full triplet