pyg-nightly 2.6.0.dev20240511__py3-none-any.whl → 2.7.0.dev20250114__py3-none-any.whl

torch_geometric/sampler/neighbor_sampler.py

@@ -2,7 +2,7 @@ import copy
 import math
 import sys
 import warnings
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable, Dict, List, Literal, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -168,7 +168,7 @@ class NeighborSampler(BaseSampler):
             attrs = [attr for attr in feature_store.get_all_tensor_attrs()]
 
             edge_attrs = graph_store.get_all_edge_attrs()
-            self.edge_types = list(set(attr.edge_type for attr in edge_attrs))
+            self.edge_types = list({attr.edge_type for attr in edge_attrs})
 
             if weight_attr is not None:
                 raise NotImplementedError(
@@ -593,7 +593,7 @@ def edge_sample(
                 src_node_time = node_time
 
             src_neg = neg_sample(src, neg_sampling, num_src_nodes, src_time,
-                                 src_node_time)
+                                 src_node_time, endpoint='src')
             src = torch.cat([src, src_neg], dim=0)
 
             if isinstance(node_time, dict):
@@ -602,7 +602,7 @@ def edge_sample(
                 dst_node_time = node_time
 
             dst_neg = neg_sample(dst, neg_sampling, num_dst_nodes, dst_time,
-                                 dst_node_time)
+                                 dst_node_time, endpoint='dst')
             dst = torch.cat([dst, dst_neg], dim=0)
 
             if edge_label is None:
@@ -623,7 +623,7 @@ def edge_sample(
                 dst_node_time = node_time
 
             dst_neg = neg_sample(dst, neg_sampling, num_dst_nodes, dst_time,
-                                 dst_node_time)
+                                 dst_node_time, endpoint='dst')
             dst = torch.cat([dst, dst_neg], dim=0)
 
             assert edge_label is None
@@ -631,7 +631,7 @@ def edge_sample(
             if edge_label_time is not None:
                 dst_time = edge_label_time.repeat(1 + neg_sampling.amount)
 
-    # Heterogeneus Neighborhood Sampling ######################################
+    # Heterogeneous Neighborhood Sampling #####################################
 
     if input_type is not None:
         seed_time_dict = None
@@ -724,7 +724,7 @@ def edge_sample(
                 src_time,
             )
 
-    # Homogeneus Neighborhood Sampling ########################################
+    # Homogeneous Neighborhood Sampling #######################################
 
     else:
 
@@ -781,12 +781,13 @@ def neg_sample(
     num_nodes: int,
     seed_time: Optional[Tensor],
     node_time: Optional[Tensor],
+    endpoint: Literal['str', 'dst'],
 ) -> Tensor:
     num_neg = math.ceil(seed.numel() * neg_sampling.amount)
 
     # TODO: Do not sample false negatives.
     if node_time is None:
-        return neg_sampling.sample(num_neg, num_nodes)
+        return neg_sampling.sample(num_neg, endpoint, num_nodes)
 
     # If we are in a temporal-sampling scenario, we need to respect the
     # timestamp of the given nodes we can use as negative examples.
@@ -800,7 +801,7 @@ def neg_sample(
     num_samples = math.ceil(neg_sampling.amount)
     seed_time = seed_time.view(1, -1).expand(num_samples, -1)
 
-    out = neg_sampling.sample(num_samples * seed.numel(), num_nodes)
+    out = neg_sampling.sample(num_samples * seed.numel(), endpoint, num_nodes)
     out = out.view(num_samples, seed.numel())
     mask = node_time[out] > seed_time  # holds all invalid samples.
     neg_sampling_complete = False
@@ -811,7 +812,7 @@ def neg_sample(
             break
 
         # Greedily search for alternative negatives.
-        out[mask] = tmp = neg_sampling.sample(num_invalid, num_nodes)
+        out[mask] = tmp = neg_sampling.sample(num_invalid, endpoint, num_nodes)
         mask[mask.clone()] = node_time[tmp] >= seed_time[mask]
 
     if not neg_sampling_complete:  # pragma: no cover

torch_geometric/testing/decorators.py

@@ -7,7 +7,9 @@ from typing import Callable
 
 import torch
 from packaging.requirements import Requirement
+from packaging.version import Version
 
+import torch_geometric
 from torch_geometric.typing import WITH_METIS, WITH_PYG_LIB, WITH_TORCH_SPARSE
 from torch_geometric.visualization.graph import has_graphviz
 
@@ -112,13 +114,8 @@ def onlyCUDA(func: Callable) -> Callable:
 def onlyXPU(func: Callable) -> Callable:
     r"""A decorator to skip tests if XPU is not found."""
     import pytest
-    try:
-        import intel_extension_for_pytorch as ipex
-        xpu_available = ipex.xpu.is_available()
-    except ImportError:
-        xpu_available = False
     return pytest.mark.skipif(
-        not xpu_available,
+        not torch_geometric.is_xpu_available(),
         reason="XPU not available",
     )(func)
 
@@ -176,24 +173,23 @@ def has_package(package: str) -> bool:
     req = Requirement(package)
     if find_spec(req.name) is None:
         return False
-    module = import_module(req.name)
-    if not hasattr(module, '__version__'):
-        return True
 
-    version = module.__version__
-    # `req.specifier` does not support `.dev` suffixes, e.g., for
-    # `pyg_lib==0.1.0.dev*`, so we manually drop them:
-    if '.dev' in version:
-        version = '.'.join(version.split('.dev')[:-1])
+    try:
+        module = import_module(req.name)
+        if not hasattr(module, '__version__'):
+            return True
 
-    return version in req.specifier
+        version = Version(module.__version__).base_version
+        return version in req.specifier
+    except Exception:
+        return False
 
 
 def withPackage(*args: str) -> Callable:
     r"""A decorator to skip tests if certain packages are not installed.
     Also supports version specification.
     """
-    na_packages = set(package for package in args if not has_package(package))
+    na_packages = {package for package in args if not has_package(package)}
 
     if len(na_packages) == 1:
         reason = f"Package {list(na_packages)[0]} not found"
@@ -227,12 +223,11 @@ def withDevice(func: Callable) -> Callable:
     if torch.cuda.is_available():
         devices.append(pytest.param(torch.device('cuda:0'), id='cuda:0'))
 
-    if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
-        try:  # Github CI may not have access to MPS hardware. Confirm:
-            torch.empty(1, device='mps')
-            devices.append(pytest.param(torch.device('mps:0'), id='mps'))
-        except RuntimeError:
-            pass
+    if torch_geometric.is_mps_available():
+        devices.append(pytest.param(torch.device('mps:0'), id='mps'))
+
+    if torch_geometric.is_xpu_available():
+        devices.append(pytest.param(torch.device('xpu:0'), id='xpu'))
 
     # Additional devices can be registered through environment variables:
     device = os.getenv('TORCH_DEVICE')

torch_geometric/transforms/__init__.py

@@ -20,6 +20,7 @@ from .target_indegree import TargetIndegree
 from .local_degree_profile import LocalDegreeProfile
 from .add_self_loops import AddSelfLoops
 from .add_remaining_self_loops import AddRemainingSelfLoops
+from .remove_self_loops import RemoveSelfLoops
 from .remove_isolated_nodes import RemoveIsolatedNodes
 from .remove_duplicated_edges import RemoveDuplicatedEdges
 from .knn_graph import KNNGraph
@@ -87,6 +88,7 @@ graph_transforms = [
     'LocalDegreeProfile',
     'AddSelfLoops',
     'AddRemainingSelfLoops',
+    'RemoveSelfLoops',
     'RemoveIsolatedNodes',
     'RemoveDuplicatedEdges',
     'KNNGraph',

torch_geometric/transforms/add_metapaths.py

@@ -37,7 +37,7 @@ class AddMetaPaths(BaseTransform):
     :class:`~torch_geometric.data.HeteroData` object as edge type
     :obj:`(src_node_type, "metapath_*", dst_node_type)`, where
     :obj:`src_node_type` and :obj:`dst_node_type` denote :math:`\mathcal{V}_1`
-    and :math:`\mathcal{V}_{\ell}`, repectively.
+    and :math:`\mathcal{V}_{\ell}`, respectively.
 
     In addition, a :obj:`metapath_dict` object is added to the
     :class:`~torch_geometric.data.HeteroData` object which maps the
@@ -108,12 +108,12 @@ class AddMetaPaths(BaseTransform):
         **kwargs: bool,
     ) -> None:
         if 'drop_orig_edges' in kwargs:
-            warnings.warn("'drop_orig_edges' is dprecated. Use "
+            warnings.warn("'drop_orig_edges' is deprecated. Use "
                           "'drop_orig_edge_types' instead")
             drop_orig_edge_types = kwargs['drop_orig_edges']
 
         if 'drop_unconnected_nodes' in kwargs:
-            warnings.warn("'drop_unconnected_nodes' is dprecated. Use "
+            warnings.warn("'drop_unconnected_nodes' is deprecated. Use "
                           "'drop_unconnected_node_types' instead")
             drop_unconnected_node_types = kwargs['drop_unconnected_nodes']
 
@@ -231,7 +231,7 @@ class AddRandomMetaPaths(BaseTransform):
             will drop node types not connected by any edge type.
             (default: :obj:`False`)
         walks_per_node (int, List[int], optional): The number of random walks
-            for each starting node in a metapth. (default: :obj:`1`)
+            for each starting node in a metapath. (default: :obj:`1`)
         sample_ratio (float, optional): The ratio of source nodes to start
             random walks from. (default: :obj:`1.0`)
     """

torch_geometric/transforms/add_positional_encoding.py

@@ -92,7 +92,7 @@ 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

torch_geometric/transforms/delaunay.py

@@ -1,4 +1,5 @@
-import scipy.spatial
+from typing import List
+
 import torch
 
 from torch_geometric.data import Data
@@ -6,28 +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 __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

@@ -2,7 +2,6 @@ from typing import Any, Dict, Tuple
 
 import numpy as np
 import torch
-from scipy.linalg import expm
 from torch import Tensor
 
 from torch_geometric.data import Data
@@ -22,7 +21,7 @@ from torch_geometric.utils import (
 @functional_transform('gdc')
 class GDC(BaseTransform):
     r"""Processes the graph via Graph Diffusion Convolution (GDC) from the
-    `"Diffusion Improves Graph Learning" <https://www.kdd.in.tum.de/gdc>`_
+    `"Diffusion Improves Graph Learning" <https://arxiv.org/abs/1911.05485>`_
     paper (functional name: :obj:`gdc`).
 
     .. note::
@@ -338,10 +337,10 @@ class GDC(BaseTransform):
 
         elif method == 'heat':
             raise NotImplementedError(
-                ('Currently no fast heat kernel is implemented. You are '
-                 'welcome to create one yourself, e.g., based on '
-                 '"Kloster and Gleich: Heat kernel based community detection '
-                 '(KDD 2014)."'))
+                'Currently no fast heat kernel is implemented. You are '
+                'welcome to create one yourself, e.g., based on '
+                '"Kloster and Gleich: Heat kernel based community detection '
+                '(KDD 2014)."')
         else:
             raise ValueError(f"Approximate GDC diffusion '{method}' unknown")
 
@@ -473,6 +472,8 @@ class GDC(BaseTransform):
 
         :rtype: (:class:`Tensor`)
         """
+        from scipy.linalg import expm
+
         if symmetric:
             e, V = torch.linalg.eigh(matrix, UPLO='U')
             diff_mat = V @ torch.diag(e.exp()) @ V.t()

torch_geometric/transforms/laplacian_lambda_max.py

@@ -1,7 +1,5 @@
 from typing import Optional
 
-from scipy.sparse.linalg import eigs, eigsh
-
 from torch_geometric.data import Data
 from torch_geometric.data.datapipes import functional_transform
 from torch_geometric.transforms import BaseTransform
@@ -41,6 +39,8 @@ class LaplacianLambdaMax(BaseTransform):
         self.is_undirected = is_undirected
 
     def forward(self, data: Data) -> Data:
+        from scipy.sparse.linalg import eigs, eigsh
+
         assert data.edge_index is not None
         num_nodes = data.num_nodes
 

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/node_property_split.py

@@ -44,7 +44,6 @@ class NodePropertySplit(BaseTransform):
             of the node property, so that nodes with greater values of the
             property are considered to be OOD (default: :obj:`True`)
 
-    Example:
     .. code-block:: python
 
         from torch_geometric.transforms import NodePropertySplit
@@ -54,7 +53,7 @@ class NodePropertySplit(BaseTransform):
 
         property_name = 'popularity'
         ratios = [0.3, 0.1, 0.1, 0.3, 0.2]
-        tranaform = NodePropertySplit(property_name, ratios)
+        transform = NodePropertySplit(property_name, ratios)
 
         data = transform(data)
     """

torch_geometric/transforms/pad.py

@@ -262,15 +262,14 @@ class Pad(BaseTransform):
     All the attributes of node types other than :obj:`v0` and :obj:`v1` are
     padded using a value of :obj:`1.0`.
     All the attributes of the :obj:`('v0', 'e0', 'v1')` edge type are padded
-    usin a value of :obj:`3.5`.
+    using a value of :obj:`3.5`.
     The :obj:`edge_attr` attributes of the :obj:`('v1', 'e0', 'v0')` edge type
     are padded using a value of :obj:`-1.5`, and any other attributes of this
     edge type are padded using a value of :obj:`5.5`.
     All the attributes of edge types other than these two are padded using a
     value of :obj:`1.5`.
 
-    Example:
-    .. code-block::
+    .. code-block:: python
 
         num_nodes = {'v0': 10, 'v1': 20, 'v2':30}
         num_edges = {('v0', 'e0', 'v1'): 80}
@@ -467,9 +466,11 @@ class Pad(BaseTransform):
         edge_type: Optional[EdgeType] = None,
     ) -> None:
 
-        attrs_to_pad = set(
-            attr for attr in store.keys()
-            if store.is_edge_attr(attr) and self.__should_pad_edge_attr(attr))
+        attrs_to_pad = {
+            attr
+            for attr in store.keys()
+            if store.is_edge_attr(attr) and self.__should_pad_edge_attr(attr)
+        }
         if not attrs_to_pad:
             return
         num_target_edges = self.max_num_edges.get_value(edge_type)

torch_geometric/transforms/random_link_split.py

@@ -23,7 +23,7 @@ class RandomLinkSplit(BaseTransform):
     in validation and test splits; and the validation split does not include
     edges in the test split.
 
-    .. code-block::
+    .. code-block:: python
 
         from torch_geometric.transforms import RandomLinkSplit
 

torch_geometric/transforms/remove_self_loops.py

@@ -0,0 +1,36 @@
+from typing import Union
+
+from torch_geometric.data import Data, HeteroData
+from torch_geometric.data.datapipes import functional_transform
+from torch_geometric.transforms import BaseTransform
+from torch_geometric.utils import remove_self_loops
+
+
+@functional_transform('remove_self_loops')
+class RemoveSelfLoops(BaseTransform):
+    r"""Removes all self-loops in the given homogeneous or heterogeneous
+    graph (functional name: :obj:`remove_self_loops`).
+
+    Args:
+        attr (str, optional): The name of the attribute of edge weights
+            or multi-dimensional edge features to pass to
+            :meth:`torch_geometric.utils.remove_self_loops`.
+            (default: :obj:`"edge_weight"`)
+    """
+    def __init__(self, attr: str = 'edge_weight') -> None:
+        self.attr = attr
+
+    def forward(
+        self,
+        data: Union[Data, HeteroData],
+    ) -> Union[Data, HeteroData]:
+        for store in data.edge_stores:
+            if store.is_bipartite() or 'edge_index' not in store:
+                continue
+
+            store.edge_index, store[self.attr] = remove_self_loops(
+                store.edge_index,
+                edge_attr=store.get(self.attr, None),
+            )
+
+        return data

torch_geometric/transforms/svd_feature_reduction.py

@@ -11,7 +11,7 @@ class SVDFeatureReduction(BaseTransform):
     Decomposition (SVD) (functional name: :obj:`svd_feature_reduction`).
 
     Args:
-        out_channels (int): The dimensionlity of node features after
+        out_channels (int): The dimensionality of node features after
             reduction.
     """
     def __init__(self, out_channels: int):

torch_geometric/transforms/virtual_node.py

@@ -37,7 +37,8 @@ class VirtualNode(BaseTransform):
         col = torch.cat([col, full, arange], dim=0)
         edge_index = torch.stack([row, col], dim=0)
 
-        new_type = edge_type.new_full((num_nodes, ), int(edge_type.max()) + 1)
+        num_edge_types = int(edge_type.max()) if edge_type.numel() > 0 else 0
+        new_type = edge_type.new_full((num_nodes, ), num_edge_types + 1)
         edge_type = torch.cat([edge_type, new_type, new_type + 1], dim=0)
 
         old_data = copy.copy(data)

torch_geometric/typing.py

@@ -13,6 +13,9 @@ WITH_PT20 = int(torch.__version__.split('.')[0]) >= 2
 WITH_PT21 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 1
 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_PT26 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 6
 WITH_PT111 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 11
 WITH_PT112 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 12
 WITH_PT113 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 13
@@ -304,6 +307,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))`:
@@ -311,27 +316,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
@@ -369,3 +384,14 @@ MaybeHeteroEdgeTensor = Union[Tensor, Dict[EdgeType, Tensor]]
 
 InputNodes = Union[OptTensor, NodeType, Tuple[NodeType, OptTensor]]
 InputEdges = Union[OptTensor, EdgeType, Tuple[EdgeType, OptTensor]]
+
+# Serialization ###############################################################
+
+if WITH_PT24:
+    torch.serialization.add_safe_globals([
+        SparseTensor,
+        SparseStorage,
+        TensorFrame,
+        MockTorchCSCTensor,
+        EdgeTypeStr,
+    ])