pyg-nightly 2.6.0.dev20240704__py3-none-any.whl → 2.8.0.dev20251207__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

@@ -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
@@ -36,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
 
@@ -87,6 +89,7 @@ graph_transforms = [
     'LocalDegreeProfile',
     'AddSelfLoops',
     'AddRemainingSelfLoops',
+    'RemoveSelfLoops',
     'RemoveIsolatedNodes',
     'RemoveDuplicatedEdges',
     'KNNGraph',
@@ -106,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

@@ -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,13 +108,15 @@ class AddMetaPaths(BaseTransform):
         **kwargs: bool,
     ) -> None:
         if 'drop_orig_edges' in kwargs:
-            warnings.warn("'drop_orig_edges' is dprecated. 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 dprecated. 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(
@@ -231,7 +233,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`)
     """
@@ -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
@@ -21,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::
@@ -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/node_property_split.py

@@ -53,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/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/pad.py

@@ -262,7 +262,7 @@ 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`.

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/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/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/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)