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

torch_geometric/loader/cluster.py

@@ -1,4 +1,5 @@
 import copy
+import os
 import os.path as osp
 import sys
 from dataclasses import dataclass
@@ -11,6 +12,7 @@ from torch import Tensor
 import torch_geometric.typing
 from torch_geometric.data import Data
 from torch_geometric.index import index2ptr, ptr2index
+from torch_geometric.io import fs
 from torch_geometric.typing import pyg_lib
 from torch_geometric.utils import index_sort, narrow, select, sort_edge_index
 from torch_geometric.utils.map import map_index
@@ -43,6 +45,8 @@ class ClusterData(torch.utils.data.Dataset):
             (default: :obj:`False`)
         save_dir (str, optional): If set, will save the partitioned data to the
             :obj:`save_dir` directory for faster re-use. (default: :obj:`None`)
+        filename (str, optional): Name of the stored partitioned file.
+            (default: :obj:`None`)
         log (bool, optional): If set to :obj:`False`, will not log any
             progress. (default: :obj:`True`)
         keep_inter_cluster_edges (bool, optional): If set to :obj:`True`,
@@ -56,6 +60,7 @@ class ClusterData(torch.utils.data.Dataset):
         num_parts: int,
         recursive: bool = False,
         save_dir: Optional[str] = None,
+        filename: Optional[str] = None,
         log: bool = True,
         keep_inter_cluster_edges: bool = False,
         sparse_format: Literal['csr', 'csc'] = 'csr',
@@ -69,11 +74,11 @@ class ClusterData(torch.utils.data.Dataset):
         self.sparse_format = sparse_format
 
         recursive_str = '_recursive' if recursive else ''
-        filename = f'metis_{num_parts}{recursive_str}.pt'
-        path = osp.join(save_dir or '', filename)
+        root_dir = osp.join(save_dir or '', f'part_{num_parts}{recursive_str}')
+        path = osp.join(root_dir, filename or 'metis.pt')
 
         if save_dir is not None and osp.exists(path):
-            self.partition = torch.load(path)
+            self.partition = fs.torch_load(path)
         else:
             if log:  # pragma: no cover
                 print('Computing METIS partitioning...', file=sys.stderr)
@@ -82,6 +87,7 @@ class ClusterData(torch.utils.data.Dataset):
             self.partition = self._partition(data.edge_index, cluster)
 
             if save_dir is not None:
+                os.makedirs(root_dir, exist_ok=True)
                 torch.save(self.partition, path)
 
             if log:  # pragma: no cover

torch_geometric/loader/graph_saint.py

@@ -4,6 +4,7 @@ from typing import Optional
 import torch
 from tqdm import tqdm
 
+from torch_geometric.io import fs
 from torch_geometric.typing import SparseTensor
 
 
@@ -77,7 +78,7 @@ class GraphSAINTSampler(torch.utils.data.DataLoader):
         if self.sample_coverage > 0:
             path = osp.join(save_dir or '', self._filename)
             if save_dir is not None and osp.exists(path):  # pragma: no cover
-                self.node_norm, self.edge_norm = torch.load(path)
+                self.node_norm, self.edge_norm = fs.torch_load(path)
             else:
                 self.node_norm, self.edge_norm = self._compute_norm()
                 if save_dir is not None:  # pragma: no cover

torch_geometric/loader/ibmb_loader.py

@@ -1,9 +1,17 @@
 import logging
 import math
-from typing import Callable, Iterator, List, NamedTuple, Optional, Tuple, Union
+from typing import (
+    Any,
+    Callable,
+    Iterator,
+    List,
+    NamedTuple,
+    Optional,
+    Tuple,
+    Union,
+)
 
 import numpy as np
-import scipy.sparse
 import torch
 from torch import Tensor
 from tqdm import tqdm
@@ -281,7 +289,7 @@ def create_batchwise_out_aux_pairs(
     return loader
 
 
-def get_pairs(ppr_mat: scipy.sparse.csr_matrix) -> np.ndarray:
+def get_pairs(ppr_mat: Any) -> np.ndarray:
     ppr_mat = ppr_mat + ppr_mat.transpose()
 
     ppr_mat = ppr_mat.tocoo()
@@ -387,7 +395,7 @@ def topk_ppr_matrix(
     output_node_indices: Union[np.ndarray, torch.LongTensor],
     topk: int,
     normalization='row',
-) -> Tuple[scipy.sparse.csr_matrix, List[np.ndarray]]:
+) -> Tuple[Any, List[np.ndarray]]:
     neighbors, weights = get_ppr(edge_index, alpha, eps, output_node_indices,
                                  num_nodes)
 

torch_geometric/loader/mixin.py

@@ -56,7 +56,7 @@ def get_numa_nodes_cores() -> Dict[str, Any]:
             nodes[numa_node_id] = sorted([(k, sorted(v))
                                           for k, v in thread_siblings.items()])
 
-    except (OSError, ValueError, IndexError, IOError):
+    except (OSError, ValueError, IndexError):
         Warning('Failed to read NUMA info')
         return {}
 

torch_geometric/loader/neighbor_loader.py

@@ -14,7 +14,7 @@ class NeighborLoader(NodeLoader):
     This loader allows for mini-batch training of GNNs on large-scale graphs
     where full-batch training is not feasible.
 
-    More specifically, :obj:`num_neighbors` denotes how much neighbors are
+    More specifically, :obj:`num_neighbors` denotes how many neighbors are
     sampled for each node in each iteration.
     :class:`~torch_geometric.loader.NeighborLoader` takes in this list of
     :obj:`num_neighbors` and iteratively samples :obj:`num_neighbors[i]` for

torch_geometric/loader/neighbor_sampler.py

@@ -72,9 +72,9 @@ class NeighborSampler(torch.utils.data.DataLoader):
         `examples/reddit.py
         <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
         reddit.py>`_ or
-        `examples/ogbn_products_sage.py
+        `examples/ogbn_train.py
         <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
-        ogbn_products_sage.py>`_.
+        ogbn_train.py>`_.
 
     Args:
         edge_index (Tensor or SparseTensor): A :obj:`torch.LongTensor` or a

torch_geometric/loader/prefetch.py

@@ -73,7 +73,7 @@ class PrefetchLoader:
         if isinstance(batch, dict):
             return {k: self.non_blocking_transfer(v) for k, v in batch.items()}
 
-        batch = batch.pin_memory(self.device_helper.device)
+        batch = batch.pin_memory()
         return batch.to(self.device_helper.device, non_blocking=True)
 
     def __iter__(self) -> Any:

torch_geometric/loader/rag_loader.py

@@ -0,0 +1,107 @@
+from abc import abstractmethod
+from typing import Any, Callable, Dict, Optional, Protocol, Tuple, Union
+
+from torch_geometric.data import Data, FeatureStore, HeteroData
+from torch_geometric.sampler import HeteroSamplerOutput, SamplerOutput
+from torch_geometric.typing import InputEdges, InputNodes
+
+
+class RAGFeatureStore(Protocol):
+    """Feature store template for remote GNN RAG backend."""
+    @abstractmethod
+    def retrieve_seed_nodes(self, query: Any, **kwargs) -> InputNodes:
+        """Makes a comparison between the query and all the nodes to get all
+        the closest nodes. Return the indices of the nodes that are to be seeds
+        for the RAG Sampler.
+        """
+        ...
+
+    @abstractmethod
+    def retrieve_seed_edges(self, query: Any, **kwargs) -> InputEdges:
+        """Makes a comparison between the query and all the edges to get all
+        the closest nodes. Returns the edge indices that are to be the seeds
+        for the RAG Sampler.
+        """
+        ...
+
+    @abstractmethod
+    def load_subgraph(
+        self, sample: Union[SamplerOutput, HeteroSamplerOutput]
+    ) -> Union[Data, HeteroData]:
+        """Combines sampled subgraph output with features in a Data object."""
+        ...
+
+
+class RAGGraphStore(Protocol):
+    """Graph store template for remote GNN RAG backend."""
+    @abstractmethod
+    def sample_subgraph(self, seed_nodes: InputNodes, seed_edges: InputEdges,
+                        **kwargs) -> Union[SamplerOutput, HeteroSamplerOutput]:
+        """Sample a subgraph using the seeded nodes and edges."""
+        ...
+
+    @abstractmethod
+    def register_feature_store(self, feature_store: FeatureStore):
+        """Register a feature store to be used with the sampler. Samplers need
+        info from the feature store in order to work properly on HeteroGraphs.
+        """
+        ...
+
+
+# TODO: Make compatible with Heterographs
+
+
+class RAGQueryLoader:
+    """Loader meant for making RAG queries from a remote backend."""
+    def __init__(self, data: Tuple[RAGFeatureStore, RAGGraphStore],
+                 local_filter: Optional[Callable[[Data, Any], Data]] = None,
+                 seed_nodes_kwargs: Optional[Dict[str, Any]] = None,
+                 seed_edges_kwargs: Optional[Dict[str, Any]] = None,
+                 sampler_kwargs: Optional[Dict[str, Any]] = None,
+                 loader_kwargs: Optional[Dict[str, Any]] = None):
+        """Loader meant for making queries from a remote backend.
+
+        Args:
+            data (Tuple[RAGFeatureStore, RAGGraphStore]): Remote FeatureStore
+                and GraphStore to load from. Assumed to conform to the
+                protocols listed above.
+            local_filter (Optional[Callable[[Data, Any], Data]], optional):
+                Optional local transform to apply to data after retrieval.
+                Defaults to None.
+            seed_nodes_kwargs (Optional[Dict[str, Any]], optional): Paramaters
+                to pass into process for fetching seed nodes. Defaults to None.
+            seed_edges_kwargs (Optional[Dict[str, Any]], optional): Parameters
+                to pass into process for fetching seed edges. Defaults to None.
+            sampler_kwargs (Optional[Dict[str, Any]], optional): Parameters to
+                pass into process for sampling graph. Defaults to None.
+            loader_kwargs (Optional[Dict[str, Any]], optional): Parameters to
+                pass into process for loading graph features. Defaults to None.
+        """
+        fstore, gstore = data
+        self.feature_store = fstore
+        self.graph_store = gstore
+        self.graph_store.register_feature_store(self.feature_store)
+        self.local_filter = local_filter
+        self.seed_nodes_kwargs = seed_nodes_kwargs or {}
+        self.seed_edges_kwargs = seed_edges_kwargs or {}
+        self.sampler_kwargs = sampler_kwargs or {}
+        self.loader_kwargs = loader_kwargs or {}
+
+    def query(self, query: Any) -> Data:
+        """Retrieve a subgraph associated with the query with all its feature
+        attributes.
+        """
+        seed_nodes = self.feature_store.retrieve_seed_nodes(
+            query, **self.seed_nodes_kwargs)
+        seed_edges = self.feature_store.retrieve_seed_edges(
+            query, **self.seed_edges_kwargs)
+
+        subgraph_sample = self.graph_store.sample_subgraph(
+            seed_nodes, seed_edges, **self.sampler_kwargs)
+
+        data = self.feature_store.load_subgraph(sample=subgraph_sample,
+                                                **self.loader_kwargs)
+
+        if self.local_filter:
+            data = self.local_filter(data, query)
+        return data

torch_geometric/loader/zip_loader.py

@@ -59,6 +59,16 @@ class ZipLoader(torch.utils.data.DataLoader):
         self.loaders = loaders
         self.filter_per_worker = filter_per_worker
 
+    def __call__(
+        self,
+        index: Union[Tensor, List[int]],
+    ) -> Union[Tuple[Data, ...], Tuple[HeteroData, ...]]:
+        r"""Samples subgraphs from a batch of input IDs."""
+        out = self.collate_fn(index)
+        if not self.filter_per_worker:
+            out = self.filter_fn(out)
+        return out
+
     def collate_fn(self, index: List[int]) -> Tuple[Any, ...]:
         if not isinstance(index, Tensor):
             index = torch.tensor(index, dtype=torch.long)

torch_geometric/metrics/__init__.py

@@ -1,14 +1,23 @@
 # flake8: noqa
 
-from .link_pred import (LinkPredPrecision, LinkPredRecall, LinkPredF1,
-                        LinkPredMAP, LinkPredNDCG)
+from .link_pred import (
+    LinkPredMetricCollection,
+    LinkPredPrecision,
+    LinkPredRecall,
+    LinkPredF1,
+    LinkPredMAP,
+    LinkPredNDCG,
+    LinkPredMRR,
+)
 
 link_pred_metrics = [
+    'LinkPredMetricCollection',
     'LinkPredPrecision',
     'LinkPredRecall',
     'LinkPredF1',
     'LinkPredMAP',
     'LinkPredNDCG',
+    'LinkPredMRR',
 ]
 
 __all__ = link_pred_metrics

torch_geometric/metrics/link_pred.py

@@ -1,4 +1,4 @@
-from typing import Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -43,34 +43,15 @@ class LinkPredMetric(BaseMetric):
             self.register_buffer('accum', torch.tensor(0.))
             self.register_buffer('total', torch.tensor(0))
 
-    def update(
-        self,
+    @staticmethod
+    def _prepare(
         pred_index_mat: Tensor,
         edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
-    ) -> None:
-        r"""Updates the state variables based on the current mini-batch
-        prediction.
-
-        :meth:`update` can be repeated multiple times to accumulate the results
-        of successive predictions, *e.g.*, inside a mini-batch training or
-        evaluation loop.
-
-        Args:
-            pred_index_mat (torch.Tensor): The top-:math:`k` predictions of
-                every example in the mini-batch with shape
-                :obj:`[batch_size, k]`.
-            edge_label_index (torch.Tensor): The ground-truth indices for every
-                example in the mini-batch, given in COO format of shape
-                :obj:`[2, num_ground_truth_indices]`.
-        """
-        if pred_index_mat.size(1) != self.k:
-            raise ValueError(f"Expected 'pred_index_mat' to hold {self.k} "
-                             f"many indices for every entry "
-                             f"(got {pred_index_mat.size(1)})")
-
-        # Compute a boolean matrix indicating if the k-th prediction is part of
-        # the ground-truth. We do this by flattening both prediction and
-        # target indices, and then determining overlaps via `torch.isin`.
+    ) -> Tuple[Tensor, Tensor]:
+        # Compute a boolean matrix indicating if the `k`-th prediction is part
+        # of the ground-truth, as well as the number of ground-truths for every
+        # example. We do this by flattening both prediction and ground-truth
+        # indices, and then determining overlaps via `torch.isin`.
         max_index = max(  # type: ignore
             pred_index_mat.max() if pred_index_mat.numel() > 0 else 0,
             edge_label_index[1].max()
@@ -78,8 +59,8 @@ class LinkPredMetric(BaseMetric):
         ) + 1
         arange = torch.arange(
             start=0,
-            end=max_index * pred_index_mat.size(0),
-            step=max_index,
+            end=max_index * pred_index_mat.size(0),  # type: ignore
+            step=max_index,  # type: ignore
             device=pred_index_mat.device,
         ).view(-1, 1)
         flat_pred_index = (pred_index_mat + arange).view(-1)
@@ -88,7 +69,7 @@ class LinkPredMetric(BaseMetric):
         pred_isin_mat = torch.isin(flat_pred_index, flat_y_index)
         pred_isin_mat = pred_isin_mat.view(pred_index_mat.size())
 
-        # Compute the number of targets per example:
+        # Compute the number of ground-truths per example:
         y_count = scatter(
             torch.ones_like(edge_label_index[0]),
             edge_label_index[0],
@@ -97,11 +78,41 @@ class LinkPredMetric(BaseMetric):
             reduce='sum',
         )
 
-        metric = self._compute(pred_isin_mat, y_count)
+        return pred_isin_mat, y_count
 
+    def _update_from_prepared(
+        self,
+        pred_isin_mat: Tensor,
+        y_count: Tensor,
+    ) -> None:
+        metric = self._compute(pred_isin_mat[:, :self.k], y_count)
         self.accum += metric.sum()
         self.total += (y_count > 0).sum()
 
+    def update(
+        self,
+        pred_index_mat: Tensor,
+        edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+    ) -> None:
+        r"""Updates the state variables based on the current mini-batch
+        prediction.
+
+        :meth:`update` can be repeated multiple times to accumulate the results
+        of successive predictions, *e.g.*, inside a mini-batch training or
+        evaluation loop.
+
+        Args:
+            pred_index_mat (torch.Tensor): The top-:math:`k` predictions of
+                every example in the mini-batch with shape
+                :obj:`[batch_size, k]`.
+            edge_label_index (torch.Tensor): The ground-truth indices for every
+                example in the mini-batch, given in COO format of shape
+                :obj:`[2, num_ground_truth_indices]`.
+        """
+        pred_isin_mat, y_count = self._prepare(pred_index_mat,
+                                               edge_label_index)
+        self._update_from_prepared(pred_isin_mat, y_count)
+
     def compute(self) -> Tensor:
         r"""Computes the final metric value."""
         if self.total == 0:
@@ -133,6 +144,103 @@ class LinkPredMetric(BaseMetric):
         return f'{self.__class__.__name__}(k={self.k})'
 
 
+class LinkPredMetricCollection(torch.nn.ModuleDict):
+    r"""A collection of metrics to reduce and speed-up computation of link
+    prediction metrics.
+
+    .. code-block:: python
+
+        from torch_geometric.metrics import (
+            LinkPredMAP,
+            LinkPredMetricCollection,
+            LinkPredPrecision,
+            LinkPredRecall,
+        )
+
+        metrics = LinkPredMetricCollection([
+            LinkPredMAP(k=10),
+            LinkPredPrecision(k=100),
+            LinkPredRecall(k=50),
+        ])
+
+        metrics.update(pred_index_mat, edge_label_index)
+        out = metrics.compute()
+        metrics.reset()
+
+        print(out)
+        >>> {'LinkPredMAP@10': tensor(0.375),
+        ...  'LinkPredPrecision@100': tensor(0.127),
+        ...  'LinkPredRecall@50': tensor(0.483)}
+
+    Args:
+        metrics: The link prediction metrics.
+    """
+    def __init__(
+        self,
+        metrics: Union[
+            List[LinkPredMetric],
+            Dict[str, LinkPredMetric],
+        ],
+    ) -> None:
+        super().__init__()
+
+        if isinstance(metrics, (list, tuple)):
+            metrics = {
+                f'{metric.__class__.__name__}@{metric.k}': metric
+                for metric in metrics
+            }
+        assert len(metrics) > 0
+        assert isinstance(metrics, dict)
+
+        for name, metric in metrics.items():
+            self[name] = metric
+
+    @property
+    def max_k(self) -> int:
+        r"""The maximum number of top-:math:`k` predictions to evaluate
+        against.
+        """
+        return max([metric.k for metric in self.values()])
+
+    def update(  # type: ignore
+        self,
+        pred_index_mat: Tensor,
+        edge_label_index: Union[Tensor, Tuple[Tensor, Tensor]],
+    ) -> None:
+        r"""Updates the state variables based on the current mini-batch
+        prediction.
+
+        :meth:`update` can be repeated multiple times to accumulate the results
+        of successive predictions, *e.g.*, inside a mini-batch training or
+        evaluation loop.
+
+        Args:
+            pred_index_mat (torch.Tensor): The top-:math:`k` predictions of
+                every example in the mini-batch with shape
+                :obj:`[batch_size, k]`.
+            edge_label_index (torch.Tensor): The ground-truth indices for every
+                example in the mini-batch, given in COO format of shape
+                :obj:`[2, num_ground_truth_indices]`.
+        """
+        pred_isin_mat, y_count = LinkPredMetric._prepare(
+            pred_index_mat, edge_label_index)
+        for metric in self.values():
+            metric._update_from_prepared(pred_isin_mat, y_count)
+
+    def compute(self) -> Dict[str, Tensor]:
+        r"""Computes the final metric values."""
+        return {name: metric.compute() for name, metric in self.items()}
+
+    def reset(self) -> None:
+        r"""Reset metric state variables to their default value."""
+        for metric in self.values():
+            metric.reset()
+
+    def __repr__(self) -> str:
+        names = [f'  {name}: {metric},\n' for name, metric in self.items()]
+        return f'{self.__class__.__name__}([\n{"".join(names)}])'
+
+
 class LinkPredPrecision(LinkPredMetric):
     r"""A link prediction metric to compute Precision @ :math:`k`.
 
@@ -182,8 +290,9 @@ class LinkPredMAP(LinkPredMetric):
     higher_is_better: bool = True
 
     def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        cum_precision = (torch.cumsum(pred_isin_mat, dim=1) /
-                         torch.arange(1, self.k + 1, device=y_count.device))
+        device = pred_isin_mat.device
+        arange = torch.arange(1, pred_isin_mat.size(1) + 1, device=device)
+        cum_precision = pred_isin_mat.cumsum(dim=1) / arange
         return ((cum_precision * pred_isin_mat).sum(dim=-1) /
                 y_count.clamp(min=1e-7, max=self.k))
 
@@ -210,9 +319,25 @@ class LinkPredNDCG(LinkPredMetric):
         self.register_buffer('idcg', cumsum(multiplier))
 
     def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
-        dcg = (pred_isin_mat * self.multiplier.view(1, -1)).sum(dim=-1)
+        multiplier = self.multiplier[:pred_isin_mat.size(1)].view(1, -1)
+        dcg = (pred_isin_mat * multiplier).sum(dim=-1)
         idcg = self.idcg[y_count.clamp(max=self.k)]
 
         out = dcg / idcg
         out[out.isnan() | out.isinf()] = 0.0
         return out
+
+
+class LinkPredMRR(LinkPredMetric):
+    r"""A link prediction metric to compute the MRR @ :math:`k` (Mean
+    Reciprocal Rank).
+
+    Args:
+        k (int): The number of top-:math:`k` predictions to evaluate against.
+    """
+    higher_is_better: bool = True
+
+    def _compute(self, pred_isin_mat: Tensor, y_count: Tensor) -> Tensor:
+        device = pred_isin_mat.device
+        arange = torch.arange(1, pred_isin_mat.size(1) + 1, device=device)
+        return (pred_isin_mat / arange).max(dim=-1)[0]

torch_geometric/nn/aggr/__init__.py

@@ -25,6 +25,7 @@ from .deep_sets import DeepSetsAggregation
 from .set_transformer import SetTransformerAggregation
 from .lcm import LCMAggregation
 from .variance_preserving import VariancePreservingAggregation
+from .patch_transformer import PatchTransformerAggregation
 
 __all__ = classes = [
     'Aggregation',
@@ -53,4 +54,5 @@ __all__ = classes = [
     'SetTransformerAggregation',
     'LCMAggregation',
     'VariancePreservingAggregation',
+    'PatchTransformerAggregation',
 ]

torch_geometric/nn/aggr/attention.py

@@ -65,8 +65,6 @@ class AttentionalAggregation(Aggregation):
                 ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
                 dim: int = -2) -> Tensor:
 
-        self.assert_two_dimensional_input(x, dim)
-
         if self.gate_mlp is not None:
             gate = self.gate_mlp(x, batch=index, batch_size=dim_size)
         else:

torch_geometric/nn/aggr/base.py

@@ -25,7 +25,7 @@ class Aggregation(torch.nn.Module):
     Notably, :obj:`index` does not have to be sorted (for most aggregation
     operators):
 
-    .. code-block::
+    .. code-block:: python
 
        # Feature matrix holding 10 elements with 64 features each:
        x = torch.randn(10, 64)
@@ -39,7 +39,7 @@ class Aggregation(torch.nn.Module):
     called :obj:`ptr`. Here, elements within the same set need to be grouped
     together in the input, and :obj:`ptr` defines their boundaries:
 
-    .. code-block::
+    .. code-block:: python
 
        # Feature matrix holding 10 elements with 64 features each:
        x = torch.randn(10, 64)
@@ -94,11 +94,9 @@ class Aggregation(torch.nn.Module):
             max_num_elements: (int, optional): The maximum number of elements
                 within a single aggregation group. (default: :obj:`None`)
         """
-        pass
 
     def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""
-        pass
 
     @disable_dynamic_shapes(required_args=['dim_size'])
     def __call__(