pyg-nightly 2.6.0.dev20240318__py3-none-any.whl → 2.7.0.dev20250115__py3-none-any.whl

torch_geometric/sampler/base.py

@@ -5,7 +5,7 @@ from abc import ABC
 from collections import defaultdict
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Union
 
 import torch
 from torch import Tensor
@@ -425,6 +425,14 @@ class NumNeighbors:
             else:
                 assert False
 
+            # Confirm that `values` only hold valid edge types:
+            if isinstance(self.values, dict):
+                edge_types_str = {EdgeTypeStr(key) for key in edge_types}
+                invalid_edge_types = set(self.values.keys()) - edge_types_str
+                if len(invalid_edge_types) > 0:
+                    raise ValueError("Not all edge types specified in "
+                                     "'num_neighbors' exist in the graph")
+
             out = {}
             for edge_type in edge_types:
                 edge_type_str = EdgeTypeStr(edge_type)
@@ -444,7 +452,7 @@ class NumNeighbors:
             out = copy.copy(self.values)
 
         if isinstance(out, dict):
-            num_hops = set(len(v) for v in out.values())
+            num_hops = {len(v) for v in out.values()}
             if len(num_hops) > 1:
                 raise ValueError(f"Number of hops must be the same across all "
                                  f"edge types (got {len(num_hops)} different "
@@ -533,24 +541,31 @@ class NegativeSampling(CastMixin):
             destination nodes for each positive source node.
         amount (int or float, optional): The ratio of sampled negative edges to
             the number of positive edges. (default: :obj:`1`)
-        weight (torch.Tensor, optional): A node-level vector determining the
-            sampling of nodes. Does not necessariyl need to sum up to one.
-            If not given, negative nodes will be sampled uniformly.
+        src_weight (torch.Tensor, optional): A node-level vector determining
+            the sampling of source nodes. Does not necessarily need to sum up
+            to one. If not given, negative nodes will be sampled uniformly.
+            (default: :obj:`None`)
+        dst_weight (torch.Tensor, optional): A node-level vector determining
+            the sampling of destination nodes. Does not necessarily need to sum
+            up to one. If not given, negative nodes will be sampled uniformly.
             (default: :obj:`None`)
     """
     mode: NegativeSamplingMode
     amount: Union[int, float] = 1
-    weight: Optional[Tensor] = None
+    src_weight: Optional[Tensor] = None
+    dst_weight: Optional[Tensor] = None
 
     def __init__(
         self,
         mode: Union[NegativeSamplingMode, str],
         amount: Union[int, float] = 1,
-        weight: Optional[Tensor] = None,
+        src_weight: Optional[Tensor] = None,
+        dst_weight: Optional[Tensor] = None,
     ):
         self.mode = NegativeSamplingMode(mode)
         self.amount = amount
-        self.weight = weight
+        self.src_weight = src_weight
+        self.dst_weight = dst_weight
 
         if self.amount <= 0:
             raise ValueError(f"The attribute 'amount' needs to be positive "
@@ -571,22 +586,28 @@ class NegativeSampling(CastMixin):
     def is_triplet(self) -> bool:
         return self.mode == NegativeSamplingMode.triplet
 
-    def sample(self, num_samples: int,
-               num_nodes: Optional[int] = None) -> Tensor:
+    def sample(
+        self,
+        num_samples: int,
+        endpoint: Literal['src', 'dst'],
+        num_nodes: Optional[int] = None,
+    ) -> Tensor:
         r"""Generates :obj:`num_samples` negative samples."""
-        if self.weight is None:
+        weight = self.src_weight if endpoint == 'src' else self.dst_weight
+
+        if weight is None:
             if num_nodes is None:
                 raise ValueError(
                     f"Cannot sample negatives in '{self.__class__.__name__}' "
                     f"without passing the 'num_nodes' argument")
             return torch.randint(num_nodes, (num_samples, ))
 
-        if num_nodes is not None and self.weight.numel() != num_nodes:
+        if num_nodes is not None and weight.numel() != num_nodes:
             raise ValueError(
                 f"The 'weight' attribute in '{self.__class__.__name__}' "
                 f"needs to match the number of nodes {num_nodes} "
                 f"(got {self.weight.numel()})")
-        return torch.multinomial(self.weight, num_samples, replacement=True)
+        return torch.multinomial(weight, num_samples, replacement=True)
 
 
 class BaseSampler(ABC):

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/sampler/utils.py

@@ -5,9 +5,9 @@ from torch import Tensor
 
 from torch_geometric.data import Data, HeteroData
 from torch_geometric.data.storage import EdgeStorage
+from torch_geometric.index import index2ptr
 from torch_geometric.typing import EdgeType, NodeType, OptTensor
 from torch_geometric.utils import coalesce, index_sort, lexsort
-from torch_geometric.utils.sparse import index2ptr
 
 # Edge Layout Conversion ######################################################
 

torch_geometric/template.py

@@ -28,6 +28,7 @@ def module_from_template(
             delete=False,
     ) as tmp:
         tmp.write(module_repr)
+        tmp.flush()
 
     spec = importlib.util.spec_from_file_location(module_name, tmp.name)
     assert spec is not None

torch_geometric/testing/__init__.py

@@ -10,7 +10,8 @@ from .decorators import (
     onlyDistributedTest,
     onlyLinux,
     noWindows,
-    onlyPython,
+    noMac,
+    minPython,
     onlyCUDA,
     onlyXPU,
     onlyOnline,
@@ -18,6 +19,7 @@ from .decorators import (
     onlyNeighborSampler,
     has_package,
     withPackage,
+    withDevice,
     withCUDA,
     withMETIS,
     disableExtensions,
@@ -39,7 +41,8 @@ __all__ = [
     'onlyDistributedTest',
     'onlyLinux',
     'noWindows',
-    'onlyPython',
+    'noMac',
+    'minPython',
     'onlyCUDA',
     'onlyXPU',
     'onlyOnline',
@@ -47,6 +50,7 @@ __all__ = [
     'onlyNeighborSampler',
     'has_package',
     'withPackage',
+    'withDevice',
     'withCUDA',
     'withMETIS',
     'disableExtensions',

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
 
@@ -67,15 +69,34 @@ def noWindows(func: Callable) -> Callable:
     )(func)
 
 
-def onlyPython(*args: str) -> Callable:
+def noMac(func: Callable) -> Callable:
+    r"""A decorator to specify that this function should not execute on
+    macOS systems.
+    """
+    import pytest
+    return pytest.mark.skipif(
+        sys.platform == 'darwin',
+        reason="macOS system",
+    )(func)
+
+
+def minPython(version: str) -> Callable:
     r"""A decorator to run tests on specific :python:`Python` versions only."""
     def decorator(func: Callable) -> Callable:
         import pytest
 
-        python_version = f'{sys.version_info.major}.{sys.version_info.minor}'
+        major, minor = version.split('.')
+
+        skip = False
+        if sys.version_info.major < int(major):
+            skip = True
+        if (sys.version_info.major == int(major)
+                and sys.version_info.minor < int(minor)):
+            skip = True
+
         return pytest.mark.skipif(
-            python_version not in args,
-            reason=f"Python {python_version} not supported",
+            skip,
+            reason=f"Python {version} required",
         )(func)
 
     return decorator
@@ -93,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)
 
@@ -157,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"
@@ -196,6 +211,24 @@ def withCUDA(func: Callable) -> Callable:
     if torch.cuda.is_available():
         devices.append(pytest.param(torch.device('cuda:0'), id='cuda:0'))
 
+    return pytest.mark.parametrize('device', devices)(func)
+
+
+def withDevice(func: Callable) -> Callable:
+    r"""A decorator to test on all available tensor processing devices."""
+    import pytest
+
+    devices = [pytest.param(torch.device('cpu'), id='cpu')]
+
+    if torch.cuda.is_available():
+        devices.append(pytest.param(torch.device('cuda:0'), id='cuda:0'))
+
+    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')
     if device:
@@ -218,10 +251,11 @@ def withMETIS(func: Callable) -> Callable:
 
     if with_metis:
         try:  # Test that METIS can succesfully execute:
-            import pyg_lib
+            # 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])
             col = torch.tensor([1, 2, 0, 2, 1, 0])
-            pyg_lib.partition.metis(rowptr, col, num_partitions=2)
+            torch.ops.torch_sparse.partition(rowptr, col, None, 2, True)
         except Exception:
             with_metis = False
 

torch_geometric/testing/feature_store.py

@@ -36,7 +36,7 @@ class MyFeatureStore(FeatureStore):
         index, tensor = self.store.get(self.key(attr), (None, None))
 
         if tensor is None:
-            return None
+            raise KeyError(f"Could not find tensor for '{attr}'")
 
         assert isinstance(tensor, Tensor)
 

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']
 
@@ -158,7 +158,7 @@ class AddMetaPaths(BaseTransform):
                         edge_index, edge_weight)
 
             new_edge_type = (metapath[0][0], f'metapath_{j}', metapath[-1][-1])
-            data[new_edge_type].edge_index = edge_index
+            data[new_edge_type].edge_index = edge_index.as_tensor()
             if self.weighted:
                 data[new_edge_type].edge_weight = edge_weight
             data.metapath_dict[new_edge_type] = metapath
@@ -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()