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

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

@@ -79,7 +79,7 @@ class ToSparseTensor(BaseTransform):
 
             keys, values = [], []
             for key, value in store.items():
-                if key == 'edge_index':
+                if key in {'edge_index', 'edge_label', 'edge_label_index'}:
                     continue
 
                 if store.is_edge_attr(key):

torch_geometric/transforms/two_hop.py

@@ -19,7 +19,7 @@ class TwoHop(BaseTransform):
 
         edge_index = EdgeIndex(edge_index, sparse_size=(N, N))
         edge_index = edge_index.sort_by('row')[0]
-        edge_index2, _ = edge_index @ edge_index
+        edge_index2 = edge_index.matmul(edge_index)[0].as_tensor()
         edge_index2, _ = remove_self_loops(edge_index2)
         edge_index = torch.cat([edge_index, edge_index2], dim=1)
 

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

@@ -1,8 +1,9 @@
 import inspect
 import os
 import sys
+import typing
 import warnings
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
 import numpy as np
 import torch
@@ -12,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
@@ -21,6 +25,16 @@ NO_MKL = 'USE_MKL=OFF' in torch.__config__.show() or WITH_WINDOWS
 
 MAX_INT64 = torch.iinfo(torch.int64).max
 
+if WITH_PT20:
+    INDEX_DTYPES: Set[torch.dtype] = {
+        torch.int32,
+        torch.int64,
+    }
+elif not typing.TYPE_CHECKING:  # pragma: no cover
+    INDEX_DTYPES: Set[torch.dtype] = {
+        torch.int64,
+    }
+
 if not hasattr(torch, 'sparse_csc'):
     torch.sparse_csc = torch.sparse_coo
 
@@ -293,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))`:
@@ -300,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
@@ -358,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,
+    ])

torch_geometric/utils/__init__.py

@@ -21,6 +21,7 @@ from ._subgraph import (get_num_hops, subgraph, k_hop_subgraph,
 from .dropout import dropout_adj, dropout_node, dropout_edge, dropout_path
 from ._homophily import homophily
 from ._assortativity import assortativity
+from ._normalize_edge_index import normalize_edge_index
 from .laplacian import get_laplacian
 from .mesh_laplacian import get_mesh_laplacian
 from .mask import mask_select, index_to_mask, mask_to_index
@@ -44,7 +45,7 @@ from .convert import to_networkit, from_networkit
 from .convert import to_trimesh, from_trimesh
 from .convert import to_cugraph, from_cugraph
 from .convert import to_dgl, from_dgl
-from .smiles import from_smiles, to_smiles
+from .smiles import from_rdmol, to_rdmol, from_smiles, to_smiles
 from .random import (erdos_renyi_graph, stochastic_blockmodel_graph,
                      barabasi_albert_graph)
 from ._negative_sampling import (negative_sampling, batched_negative_sampling,
@@ -89,6 +90,7 @@ __all__ = [
     'dropout_adj',
     'homophily',
     'assortativity',
+    'normalize_edge_index',
     'get_laplacian',
     'get_mesh_laplacian',
     'mask_select',
@@ -127,6 +129,8 @@ __all__ = [
     'from_cugraph',
     'to_dgl',
     'from_dgl',
+    'from_rdmol',
+    'to_rdmol',
     'from_smiles',
     'to_smiles',
     'erdos_renyi_graph',

torch_geometric/utils/_negative_sampling.py

@@ -265,7 +265,7 @@ def structured_negative_sampling_feasible(
     :meth:`~torch_geometric.utils.structured_negative_sampling` is feasible
     on the graph given by :obj:`edge_index`.
     :meth:`~torch_geometric.utils.structured_negative_sampling` is infeasible
-    if atleast one node is connected to all other nodes.
+    if at least one node is connected to all other nodes.
 
     Args:
         edge_index (LongTensor): The edge indices.

torch_geometric/utils/_normalize_edge_index.py

@@ -0,0 +1,46 @@
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from torch_geometric.utils import add_self_loops as add_self_loops_fn
+from torch_geometric.utils import degree
+
+
+def normalize_edge_index(
+    edge_index: Tensor,
+    num_nodes: Optional[int] = None,
+    add_self_loops: bool = True,
+    symmetric: bool = True,
+) -> Tuple[Tensor, Tensor]:
+    """Applies normalization to the edges of a graph.
+
+    This function can add self-loops to the graph and apply either symmetric or
+    asymmetric normalization based on the node degrees.
+
+    Args:
+        edge_index (LongTensor): The edge indices.
+        num_nodes (int, int], optional): The number of nodes, *i.e.*
+            :obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`)
+        add_self_loops (bool, optional): If set to :obj:`False`, will not add
+            self-loops to the input graph. (default: :obj:`True`)
+        symmetric (bool, optional):  If set to :obj:`True`, symmetric
+            normalization (:math:`D^{-1/2} A D^{-1/2}`) is used, otherwise
+            asymmetric normalization (:math:`D^{-1} A`).
+    """
+    if add_self_loops:
+        edge_index, _ = add_self_loops_fn(edge_index, num_nodes=num_nodes)
+
+    row, col = edge_index[0], edge_index[1]
+    deg = degree(row, num_nodes, dtype=torch.get_default_dtype())
+
+    if symmetric:  # D^-1/2 * A * D^-1/2
+        deg_inv_sqrt = deg.pow(-0.5)
+        deg_inv_sqrt[torch.isinf(deg_inv_sqrt)] = 0
+        edge_weight = deg_inv_sqrt[row] * deg_inv_sqrt[col]
+    else:  # D^-1 * A
+        deg_inv = deg.pow(-1)
+        deg_inv[torch.isinf(deg_inv)] = 0
+        edge_weight = deg_inv[row]
+
+    return edge_index, edge_weight

torch_geometric/utils/_scatter.py

@@ -4,7 +4,7 @@ import torch
 from torch import Tensor
 
 import torch_geometric.typing
-from torch_geometric import is_compiling, warnings
+from torch_geometric import is_compiling, is_in_onnx_export, warnings
 from torch_geometric.typing import torch_scatter
 from torch_geometric.utils.functions import cumsum
 
@@ -88,18 +88,33 @@ if torch_geometric.typing.WITH_PT112:  # pragma: no cover
         # in case the input does not require gradients:
         if reduce in ['min', 'max', 'amin', 'amax']:
             if (not torch_geometric.typing.WITH_TORCH_SCATTER
-                    or is_compiling() or not src.is_cuda
+                    or is_compiling() or is_in_onnx_export() or not src.is_cuda
                     or not src.requires_grad):
 
-                if src.is_cuda and src.requires_grad and not is_compiling():
+                if (src.is_cuda and src.requires_grad and not is_compiling()
+                        and not is_in_onnx_export()):
                     warnings.warn(f"The usage of `scatter(reduce='{reduce}')` "
                                   f"can be accelerated via the 'torch-scatter'"
                                   f" package, but it was not found")
 
                 index = broadcast(index, src, dim)
-                return src.new_zeros(size).scatter_reduce_(
-                    dim, index, src, reduce=f'a{reduce[-3:]}',
-                    include_self=False)
+                if not is_in_onnx_export():
+                    return src.new_zeros(size).scatter_reduce_(
+                        dim, index, src, reduce=f'a{reduce[-3:]}',
+                        include_self=False)
+
+                fill = torch.full(  # type: ignore
+                    size=(1, ),
+                    fill_value=src.min() if 'max' in reduce else src.max(),
+                    dtype=src.dtype,
+                    device=src.device,
+                ).expand_as(src)
+                out = src.new_zeros(size).scatter_reduce_(
+                    dim, index, fill, reduce=f'a{reduce[-3:]}',
+                    include_self=True)
+                return out.scatter_reduce_(dim, index, src,
+                                           reduce=f'a{reduce[-3:]}',
+                                           include_self=True)
 
             return torch_scatter.scatter(src, index, dim, dim_size=dim_size,
                                          reduce=reduce[-3:])
@@ -175,6 +190,7 @@ else:  # pragma: no cover
 
 
 def broadcast(src: Tensor, ref: Tensor, dim: int) -> Tensor:
+    dim = ref.dim() + dim if dim < 0 else dim
     size = ((1, ) * dim) + (-1, ) + ((1, ) * (ref.dim() - dim - 1))
     return src.view(size).expand_as(ref)
 
@@ -186,7 +202,8 @@ def scatter_argmax(
     dim_size: Optional[int] = None,
 ) -> Tensor:
 
-    if torch_geometric.typing.WITH_TORCH_SCATTER and not is_compiling():
+    if (torch_geometric.typing.WITH_TORCH_SCATTER and not is_compiling()
+            and not is_in_onnx_export()):
         out = torch_scatter.scatter_max(src, index, dim=dim, dim_size=dim_size)
         return out[1]
 
@@ -199,9 +216,18 @@ def scatter_argmax(
         dim_size = int(index.max()) + 1 if index.numel() > 0 else 0
 
     if torch_geometric.typing.WITH_PT112:
-        res = src.new_empty(dim_size)
-        res.scatter_reduce_(0, index, src.detach(), reduce='amax',
-                            include_self=False)
+        if not is_in_onnx_export():
+            res = src.new_empty(dim_size)
+            res.scatter_reduce_(0, index, src.detach(), reduce='amax',
+                                include_self=False)
+        else:
+            # `include_self=False` is currently not supported by ONNX:
+            res = src.new_full(
+                size=(dim_size, ),
+                fill_value=src.min(),  # type: ignore
+            )
+            res.scatter_reduce_(0, index, src.detach(), reduce="amax",
+                                include_self=True)
     elif torch_geometric.typing.WITH_PT111:
         res = torch.scatter_reduce(src.detach(), 0, index, reduce='amax',
                                    output_size=dim_size)  # type: ignore
@@ -294,7 +320,7 @@ def group_cat(
     r"""Concatenates the given sequence of tensors :obj:`tensors` in the given
     dimension :obj:`dim`.
     Different from :meth:`torch.cat`, values along the concatenating dimension
-    are grouped according to the indicies defined in the :obj:`index` tensors.
+    are grouped according to the indices defined in the :obj:`index` tensors.
     All tensors must have the same shape (except in the concatenating
     dimension).
 
@@ -325,5 +351,5 @@ def group_cat(
     """
     assert len(tensors) == len(indices)
     index, perm = torch.cat(indices).sort(stable=True)
-    out = torch.cat(tensors, dim=0)[perm]
+    out = torch.cat(tensors, dim=dim).index_select(dim, perm)
     return (out, index) if return_index else out

torch_geometric/utils/_subgraph.py

@@ -346,10 +346,12 @@ def k_hop_subgraph(
 
     subsets = [node_idx]
 
+    preserved_edge_mask = torch.zeros_like(edge_mask)
     for _ in range(num_hops):
         node_mask.fill_(False)
         node_mask[subsets[-1]] = True
         torch.index_select(node_mask, 0, row, out=edge_mask)
+        preserved_edge_mask |= edge_mask
         subsets.append(col[edge_mask])
 
     subset, inv = torch.cat(subsets).unique(return_inverse=True)
@@ -360,6 +362,8 @@ def k_hop_subgraph(
 
     if not directed:
         edge_mask = node_mask[row] & node_mask[col]
+    else:
+        edge_mask = preserved_edge_mask
 
     edge_index = edge_index[:, edge_mask]
 

torch_geometric/utils/_tree_decomposition.py

@@ -2,7 +2,6 @@ from itertools import chain
 from typing import Any, List, Literal, Tuple, Union, overload
 
 import torch
-from scipy.sparse.csgraph import minimum_spanning_tree
 from torch import Tensor
 
 from torch_geometric.utils import (
@@ -54,6 +53,7 @@ def tree_decomposition(
         :obj:`False`, else :obj:`(LongTensor, LongTensor, int, LongTensor)`
     """
     import rdkit.Chem as Chem
+    from scipy.sparse.csgraph import minimum_spanning_tree
 
     # Cliques = rings and bonds.
     cliques: List[List[int]] = [list(x) for x in Chem.GetSymmSSSR(mol)]
@@ -64,7 +64,7 @@ def tree_decomposition(
             xs.append(1)
 
     # Generate `atom2cliques` mappings.
-    atom2cliques: List[List[int]] = [[] for i in range(mol.GetNumAtoms())]
+    atom2cliques: List[List[int]] = [[] for _ in range(mol.GetNumAtoms())]
     for c in range(len(cliques)):
         for atom in cliques[c]:
             atom2cliques[atom].append(c)

torch_geometric/utils/augmentation.py

@@ -12,7 +12,7 @@ def shuffle_node(
     training: bool = True,
 ) -> Tuple[Tensor, Tensor]:
     r"""Randomly shuffle the feature matrix :obj:`x` along the
-    first dimmension.
+    first dimension.
 
     The method returns (1) the shuffled :obj:`x`, (2) the permutation
     indicating the orders of original nodes after shuffling.

torch_geometric/utils/convert.py

@@ -1,7 +1,6 @@
 from collections import defaultdict
 from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple, Union
 
-import scipy.sparse
 import torch
 from torch import Tensor
 from torch.utils.dlpack import from_dlpack, to_dlpack
@@ -14,7 +13,7 @@ def to_scipy_sparse_matrix(
     edge_index: Tensor,
     edge_attr: Optional[Tensor] = None,
     num_nodes: Optional[int] = None,
-) -> scipy.sparse.coo_matrix:
+) -> Any:
     r"""Converts a graph given by edge indices and edge attributes to a scipy
     sparse matrix.
 
@@ -34,22 +33,23 @@ def to_scipy_sparse_matrix(
         <4x4 sparse matrix of type '<class 'numpy.float32'>'
             with 6 stored elements in COOrdinate format>
     """
+    import scipy.sparse as sp
+
     row, col = edge_index.cpu()
 
     if edge_attr is None:
-        edge_attr = torch.ones(row.size(0))
+        edge_attr = torch.ones(row.size(0), device="cpu")
     else:
         edge_attr = edge_attr.view(-1).cpu()
         assert edge_attr.size(0) == row.size(0)
 
     N = maybe_num_nodes(edge_index, num_nodes)
-    out = scipy.sparse.coo_matrix(
+    out = sp.coo_matrix(  #
         (edge_attr.numpy(), (row.numpy(), col.numpy())), (N, N))
     return out
 
 
-def from_scipy_sparse_matrix(
-        A: scipy.sparse.spmatrix) -> Tuple[Tensor, Tensor]:
+def from_scipy_sparse_matrix(A: Any) -> Tuple[Tensor, Tensor]:
     r"""Converts a scipy sparse matrix to edge indices and edge attributes.
 
     Args:
@@ -527,10 +527,14 @@ def to_dgl(
     if isinstance(data, Data):
         if data.edge_index is not None:
             row, col = data.edge_index
-        else:
+        elif 'adj' in data:
+            row, col, _ = data.adj.coo()
+        elif 'adj_t' in data:
             row, col, _ = data.adj_t.t().coo()
+        else:
+            row, col = [], []
 
-        g = dgl.graph((row, col))
+        g = dgl.graph((row, col), num_nodes=data.num_nodes)
 
         for attr in data.node_attrs():
             g.ndata[attr] = data[attr]

torch_geometric/utils/geodesic.py

@@ -2,6 +2,7 @@ import multiprocessing as mp
 import warnings
 from typing import Optional
 
+import numpy as np
 import torch
 from torch import Tensor
 
@@ -82,54 +83,55 @@ def geodesic_distance(  # noqa: D417
 
     dtype = pos.dtype
 
-    pos = pos.detach().cpu().to(torch.double).numpy()
-    face = face.detach().t().cpu().to(torch.int).numpy()
+    pos_np = pos.detach().cpu().to(torch.double).numpy()
+    face_np = face.detach().t().cpu().to(torch.int).numpy()
 
     if src is None and dst is None:
-        out = gdist.local_gdist_matrix(pos, face,
-                                       max_distance * scale).toarray() / scale
+        out = gdist.local_gdist_matrix(
+            pos_np,
+            face_np,
+            max_distance * scale,
+        ).toarray() / scale
         return torch.from_numpy(out).to(dtype)
 
     if src is None:
-        src = torch.arange(pos.shape[0], dtype=torch.int).numpy()
+        src_np = torch.arange(pos.size(0), dtype=torch.int).numpy()
     else:
-        src = src.detach().cpu().to(torch.int).numpy()
-    assert src is not None
+        src_np = src.detach().cpu().to(torch.int).numpy()
 
-    dst = None if dst is None else dst.detach().cpu().to(torch.int).numpy()
+    dst_np = None if dst is None else dst.detach().cpu().to(torch.int).numpy()
 
     def _parallel_loop(
-        pos: Tensor,
-        face: Tensor,
-        src: Tensor,
-        dst: Optional[Tensor],
+        pos_np: np.ndarray,
+        face_np: np.ndarray,
+        src_np: np.ndarray,
+        dst_np: Optional[np.ndarray],
         max_distance: float,
         scale: float,
         i: int,
         dtype: torch.dtype,
     ) -> Tensor:
-        s = src[i:i + 1]
-        d = None if dst is None else dst[i:i + 1]
-        out = gdist.compute_gdist(pos, face, s, d, max_distance * scale)
+        s = src_np[i:i + 1]
+        d = None if dst_np is None else dst_np[i:i + 1]
+        out = gdist.compute_gdist(pos_np, face_np, s, d, max_distance * scale)
         out = out / scale
         return torch.from_numpy(out).to(dtype)
 
     num_workers = mp.cpu_count() if num_workers <= -1 else num_workers
     if num_workers > 0:
         with mp.Pool(num_workers) as pool:
-            outs = pool.starmap(
-                _parallel_loop,
-                [(pos, face, src, dst, max_distance, scale, i, dtype)
-                 for i in range(len(src))])
+            data = [(pos_np, face_np, src_np, dst_np, max_distance, scale, i,
+                     dtype) for i in range(len(src_np))]
+            outs = pool.starmap(_parallel_loop, data)
     else:
         outs = [
-            _parallel_loop(pos, face, src, dst, max_distance, scale, i, dtype)
-            for i in range(len(src))
+            _parallel_loop(pos_np, face_np, src_np, dst_np, max_distance,
+                           scale, i, dtype) for i in range(len(src_np))
         ]
 
     out = torch.cat(outs, dim=0)
 
     if dst is None:
-        out = out.view(-1, pos.shape[0])
+        out = out.view(-1, pos.size(0))
 
     return out

torch_geometric/utils/hetero.py

@@ -63,7 +63,7 @@ def group_hetero_graph(
 
 def get_unused_node_types(node_types: List[NodeType],
                           edge_types: List[EdgeType]) -> Set[NodeType]:
-    dst_node_types = set(edge_type[-1] for edge_type in edge_types)
+    dst_node_types = {edge_type[-1] for edge_type in edge_types}
     return set(node_types) - set(dst_node_types)