pyg-nightly 2.6.0.dev20240704__py3-none-any.whl → 2.8.0.dev20251207__py3-none-any.whl

torch_geometric/typing.py

@@ -3,7 +3,7 @@ import os
 import sys
 import typing
 import warnings
-from typing import Any, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Dict, List, Optional, Set, Tuple, TypeAlias, Union
 
 import numpy as np
 import torch
@@ -14,8 +14,10 @@ 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_PT111 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 11
-WITH_PT112 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 12
+WITH_PT25 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 5
+WITH_PT26 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 6
+WITH_PT27 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 7
+WITH_PT28 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 8
 WITH_PT113 = WITH_PT20 or int(torch.__version__.split('.')[1]) >= 13
 
 WITH_WINDOWS = os.name == 'nt'
@@ -62,10 +64,21 @@ try:
         pyg_lib.sampler.neighbor_sample).parameters)
     WITH_WEIGHTED_NEIGHBOR_SAMPLE = ('edge_weight' in inspect.signature(
         pyg_lib.sampler.neighbor_sample).parameters)
+    try:
+        torch.classes.pyg.CPUHashMap  # noqa: B018
+        WITH_CPU_HASH_MAP = True
+    except Exception:
+        WITH_CPU_HASH_MAP = False
+    try:
+        torch.classes.pyg.CUDAHashMap  # noqa: B018
+        WITH_CUDA_HASH_MAP = True
+    except Exception:
+        WITH_CUDA_HASH_MAP = False
 except Exception as e:
     if not isinstance(e, ImportError):  # pragma: no cover
-        warnings.warn(f"An issue occurred while importing 'pyg-lib'. "
-                      f"Disabling its usage. Stacktrace: {e}")
+        warnings.warn(
+            f"An issue occurred while importing 'pyg-lib'. "
+            f"Disabling its usage. Stacktrace: {e}", stacklevel=2)
     pyg_lib = object
     WITH_PYG_LIB = False
     WITH_GMM = False
@@ -76,14 +89,41 @@ except Exception as e:
     WITH_METIS = False
     WITH_EDGE_TIME_NEIGHBOR_SAMPLE = False
     WITH_WEIGHTED_NEIGHBOR_SAMPLE = False
+    WITH_CPU_HASH_MAP = False
+    WITH_CUDA_HASH_MAP = False
+
+if WITH_CPU_HASH_MAP:
+    CPUHashMap: TypeAlias = torch.classes.pyg.CPUHashMap  # type: ignore[name-defined]  # noqa: E501
+else:
+
+    class CPUHashMap:  # type: ignore
+        def __init__(self, key: Tensor) -> None:
+            raise ImportError("'CPUHashMap' requires 'pyg-lib'")
+
+        def get(self, query: Tensor) -> Tensor:
+            raise ImportError("'CPUHashMap' requires 'pyg-lib'")
+
+
+if WITH_CUDA_HASH_MAP:
+    CUDAHashMap: TypeAlias = torch.classes.pyg.CUDAHashMap  # type: ignore[name-defined]  # noqa: E501
+else:
+
+    class CUDAHashMap:  # type: ignore
+        def __init__(self, key: Tensor) -> None:
+            raise ImportError("'CUDAHashMap' requires 'pyg-lib'")
+
+        def get(self, query: Tensor) -> Tensor:
+            raise ImportError("'CUDAHashMap' requires 'pyg-lib'")
+
 
 try:
     import torch_scatter  # noqa
     WITH_TORCH_SCATTER = True
 except Exception as e:
     if not isinstance(e, ImportError):  # pragma: no cover
-        warnings.warn(f"An issue occurred while importing 'torch-scatter'. "
-                      f"Disabling its usage. Stacktrace: {e}")
+        warnings.warn(
+            f"An issue occurred while importing 'torch-scatter'. "
+            f"Disabling its usage. Stacktrace: {e}", stacklevel=2)
     torch_scatter = object
     WITH_TORCH_SCATTER = False
 
@@ -93,8 +133,9 @@ try:
     WITH_TORCH_CLUSTER_BATCH_SIZE = 'batch_size' in torch_cluster.knn.__doc__
 except Exception as e:
     if not isinstance(e, ImportError):  # pragma: no cover
-        warnings.warn(f"An issue occurred while importing 'torch-cluster'. "
-                      f"Disabling its usage. Stacktrace: {e}")
+        warnings.warn(
+            f"An issue occurred while importing 'torch-cluster'. "
+            f"Disabling its usage. Stacktrace: {e}", stacklevel=2)
     WITH_TORCH_CLUSTER = False
     WITH_TORCH_CLUSTER_BATCH_SIZE = False
 
@@ -111,7 +152,7 @@ except Exception as e:
     if not isinstance(e, ImportError):  # pragma: no cover
         warnings.warn(
             f"An issue occurred while importing 'torch-spline-conv'. "
-            f"Disabling its usage. Stacktrace: {e}")
+            f"Disabling its usage. Stacktrace: {e}", stacklevel=2)
     WITH_TORCH_SPLINE_CONV = False
 
 try:
@@ -120,8 +161,9 @@ try:
     WITH_TORCH_SPARSE = True
 except Exception as e:
     if not isinstance(e, ImportError):  # pragma: no cover
-        warnings.warn(f"An issue occurred while importing 'torch-sparse'. "
-                      f"Disabling its usage. Stacktrace: {e}")
+        warnings.warn(
+            f"An issue occurred while importing 'torch-sparse'. "
+            f"Disabling its usage. Stacktrace: {e}", stacklevel=2)
     WITH_TORCH_SPARSE = False
 
     class SparseStorage:  # type: ignore
@@ -305,6 +347,8 @@ class EdgeTypeStr(str):
     r"""A helper class to construct serializable edge types by merging an edge
     type tuple into a single string.
     """
+    edge_type: tuple[str, str, str]
+
     def __new__(cls, *args: Any) -> 'EdgeTypeStr':
         if isinstance(args[0], (list, tuple)):
             # Unwrap `EdgeType((src, rel, dst))` and `EdgeTypeStr((src, dst))`:
@@ -312,27 +356,37 @@ class EdgeTypeStr(str):
 
         if len(args) == 1 and isinstance(args[0], str):
             arg = args[0]  # An edge type string was passed.
+            edge_type = tuple(arg.split(EDGE_TYPE_STR_SPLIT))
+            if len(edge_type) != 3:
+                raise ValueError(f"Cannot convert the edge type '{arg}' to a "
+                                 f"tuple since it holds invalid characters")
 
         elif len(args) == 2 and all(isinstance(arg, str) for arg in args):
             # A `(src, dst)` edge type was passed - add `DEFAULT_REL`:
-            arg = EDGE_TYPE_STR_SPLIT.join((args[0], DEFAULT_REL, args[1]))
+            edge_type = (args[0], DEFAULT_REL, args[1])
+            arg = EDGE_TYPE_STR_SPLIT.join(edge_type)
 
         elif len(args) == 3 and all(isinstance(arg, str) for arg in args):
             # A `(src, rel, dst)` edge type was passed:
+            edge_type = tuple(args)
             arg = EDGE_TYPE_STR_SPLIT.join(args)
 
         else:
             raise ValueError(f"Encountered invalid edge type '{args}'")
 
-        return str.__new__(cls, arg)
+        out = str.__new__(cls, arg)
+        out.edge_type = edge_type  # type: ignore
+        return out
 
     def to_tuple(self) -> EdgeType:
         r"""Returns the original edge type."""
-        out = tuple(self.split(EDGE_TYPE_STR_SPLIT))
-        if len(out) != 3:
+        if len(self.edge_type) != 3:
             raise ValueError(f"Cannot convert the edge type '{self}' to a "
                              f"tuple since it holds invalid characters")
-        return out
+        return self.edge_type
+
+    def __reduce__(self) -> tuple[Any, Any]:
+        return (self.__class__, (self.edge_type, ))
 
 
 # There exist some short-cuts to query edge-types (given that the full triplet
@@ -370,3 +424,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
@@ -52,10 +53,11 @@ from ._negative_sampling import (negative_sampling, batched_negative_sampling,
                                  structured_negative_sampling_feasible)
 from .augmentation import shuffle_node, mask_feature, add_random_edge
 from ._tree_decomposition import tree_decomposition
-from .embedding import get_embeddings
+from .embedding import get_embeddings, get_embeddings_hetero
 from ._trim_to_layer import trim_to_layer
 from .ppr import get_ppr
 from ._train_test_split_edges import train_test_split_edges
+from .influence import total_influence
 
 __all__ = [
     'scatter',
@@ -89,6 +91,7 @@ __all__ = [
     'dropout_adj',
     'homophily',
     'assortativity',
+    'normalize_edge_index',
     'get_laplacian',
     'get_mesh_laplacian',
     'mask_select',
@@ -143,9 +146,11 @@ __all__ = [
     'add_random_edge',
     'tree_decomposition',
     'get_embeddings',
+    'get_embeddings_hetero',
     'trim_to_layer',
     'get_ppr',
     'train_test_split_edges',
+    'total_influence',
 ]
 
 # `structured_negative_sampling_feasible` is a long name and thus destroys the

torch_geometric/utils/_lexsort.py

@@ -1,11 +1,7 @@
 from typing import List
 
-import numpy as np
-import torch
 from torch import Tensor
 
-import torch_geometric.typing
-
 
 def lexsort(
     keys: List[Tensor],
@@ -28,11 +24,6 @@ def lexsort(
     """
     assert len(keys) >= 1
 
-    if not torch_geometric.typing.WITH_PT113:
-        keys = [k.neg() for k in keys] if descending else keys
-        out = np.lexsort([k.detach().cpu().numpy() for k in keys], axis=dim)
-        return torch.from_numpy(out).to(keys[0].device)
-
     out = keys[0].argsort(dim=dim, descending=descending, stable=True)
     for k in keys[1:]:
         index = k.gather(dim, out)

torch_geometric/utils/_negative_sampling.py

@@ -12,7 +12,7 @@ from torch_geometric.utils.num_nodes import maybe_num_nodes
 def negative_sampling(
     edge_index: Tensor,
     num_nodes: Optional[Union[int, Tuple[int, int]]] = None,
-    num_neg_samples: Optional[int] = None,
+    num_neg_samples: Optional[Union[int, float]] = None,
     method: str = "sparse",
     force_undirected: bool = False,
 ) -> Tensor:
@@ -25,10 +25,12 @@ def negative_sampling(
             If given as a tuple, then :obj:`edge_index` is interpreted as a
             bipartite graph with shape :obj:`(num_src_nodes, num_dst_nodes)`.
             (default: :obj:`None`)
-        num_neg_samples (int, optional): The (approximate) number of negative
-            samples to return.
-            If set to :obj:`None`, will try to return a negative edge for every
-            positive edge. (default: :obj:`None`)
+        num_neg_samples (int or float, optional): The (approximate) number of
+            negative samples to return. If set to a floating-point value, it
+            represents the ratio of negative samples to generate based on the
+            number of positive edges. If set to :obj:`None`, will try to
+            return a negative edge for every positive edge.
+            (default: :obj:`None`)
         method (str, optional): The method to use for negative sampling,
             *i.e.* :obj:`"sparse"` or :obj:`"dense"`.
             This is a memory/runtime trade-off.
@@ -48,6 +50,11 @@ def negative_sampling(
         tensor([[3, 0, 0, 3],
                 [2, 3, 2, 1]])
 
+        >>> negative_sampling(edge_index, num_nodes=(3, 4),
+        ...                   num_neg_samples=0.5)  # 50% of positive edges
+        tensor([[0, 3],
+                [3, 0]])
+
         >>> # For bipartite graph
         >>> negative_sampling(edge_index, num_nodes=(3, 4))
         tensor([[0, 2, 2, 1],
@@ -74,6 +81,8 @@ def negative_sampling(
 
     if num_neg_samples is None:
         num_neg_samples = edge_index.size(1)
+    elif isinstance(num_neg_samples, float):
+        num_neg_samples = int(num_neg_samples * edge_index.size(1))
     if force_undirected:
         num_neg_samples = num_neg_samples // 2
 
@@ -100,10 +109,9 @@ def negative_sampling(
         idx = idx.to('cpu')
         for _ in range(3):  # Number of tries to sample negative indices.
             rnd = sample(population, sample_size, device='cpu')
-            mask = np.isin(rnd.numpy(), idx.numpy())  # type: ignore
+            mask = torch.from_numpy(np.isin(rnd.numpy(), idx.numpy())).bool()
             if neg_idx is not None:
-                mask |= np.isin(rnd, neg_idx.to('cpu'))
-            mask = torch.from_numpy(mask).to(torch.bool)
+                mask |= torch.from_numpy(np.isin(rnd, neg_idx.cpu())).bool()
             rnd = rnd[~mask].to(edge_index.device)
             neg_idx = rnd if neg_idx is None else torch.cat([neg_idx, rnd])
             if neg_idx.numel() >= num_neg_samples:
@@ -117,7 +125,7 @@ def negative_sampling(
 def batched_negative_sampling(
     edge_index: Tensor,
     batch: Union[Tensor, Tuple[Tensor, Tensor]],
-    num_neg_samples: Optional[int] = None,
+    num_neg_samples: Optional[Union[int, float]] = None,
     method: str = "sparse",
     force_undirected: bool = False,
 ) -> Tensor:
@@ -131,9 +139,11 @@ def batched_negative_sampling(
             node to a specific example.
             If given as a tuple, then :obj:`edge_index` is interpreted as a
             bipartite graph connecting two different node types.
-        num_neg_samples (int, optional): The number of negative samples to
-            return. If set to :obj:`None`, will try to return a negative edge
-            for every positive edge. (default: :obj:`None`)
+        num_neg_samples (int or float, optional): The number of negative
+            samples to return. If set to :obj:`None`, will try to return a
+            negative edge for every positive edge. If float, it will generate
+            :obj:`num_neg_samples * num_edges` negative samples.
+            (default: :obj:`None`)
         method (str, optional): The method to use for negative sampling,
             *i.e.* :obj:`"sparse"` or :obj:`"dense"`.
             This is a memory/runtime trade-off.
@@ -157,6 +167,11 @@ def batched_negative_sampling(
         tensor([[3, 1, 3, 2, 7, 7, 6, 5],
                 [2, 0, 1, 1, 5, 6, 4, 4]])
 
+        >>> # Using float multiplier for negative samples
+        >>> batched_negative_sampling(edge_index, batch, num_neg_samples=1.5)
+        tensor([[3, 1, 3, 2, 7, 7, 6, 5, 2, 0, 1, 1],
+                [2, 0, 1, 1, 5, 6, 4, 4, 3, 2, 3, 0]])
+
         >>> # For bipartite graph
         >>> edge_index1 = torch.as_tensor([[0, 0, 1, 1], [0, 1, 2, 3]])
         >>> edge_index2 = edge_index1 + torch.tensor([[2], [4]])
@@ -265,7 +280,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