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

torch_geometric/inspector.py

@@ -305,7 +305,7 @@ class Inspector:
         according to its function signature from a data blob.
 
         Args:
-            func (callabel or str): The function.
+            func (callable or str): The function.
             kwargs (dict[str, Any]): The data blob which may serve as inputs.
         """
         out_dict: Dict[str, Any] = {}
@@ -346,7 +346,7 @@ class Inspector:
             type annotations are not found.
 
         Args:
-            func (callabel or str): The function.
+            func (callable or str): The function.
             exclude (list[int or str]): A list of parameters to exclude, either
                 given by their name or index. (default: :obj:`None`)
         """
@@ -401,7 +401,8 @@ class Inspector:
             match = find_parenthesis_content(source, f'self.{func_name}')
             if match is not None:
                 for i, kwarg in enumerate(split(match, sep=',')):
-                    if exclude is not None and i in exclude:
+                    if ('=' not in kwarg and exclude is not None
+                            and i in exclude):
                         continue
 
                     name_and_content = re.split(r'\s*=\s*', kwarg)
@@ -447,6 +448,10 @@ def type_repr(obj: Any, _globals: Dict[str, Any]) -> str:
         return '...'
 
     if obj.__module__ == 'typing':  # Special logic for `typing.*` types:
+
+        if not hasattr(obj, '_name'):
+            return repr(obj)
+
         name = obj._name
         if name is None:  # In some cases, `_name` is not populated.
             name = str(obj.__origin__).split('.')[-1]

torch_geometric/io/fs.py

@@ -1,6 +1,9 @@
 import io
 import os.path as osp
+import pickle
+import re
 import sys
+import warnings
 from typing import Any, Dict, List, Literal, Optional, Union, overload
 from uuid import uuid4
 
@@ -186,11 +189,11 @@ def rm(path: str, recursive: bool = True) -> None:
     get_fs(path).rm(path, recursive)
 
 
-def mv(path1: str, path2: str, recursive: bool = True) -> None:
+def mv(path1: str, path2: str) -> None:
     fs1 = get_fs(path1)
     fs2 = get_fs(path2)
     assert fs1.protocol == fs2.protocol
-    fs1.mv(path1, path2, recursive)
+    fs1.mv(path1, path2)
 
 
 def glob(path: str) -> List[str]:
@@ -211,5 +214,28 @@ def torch_save(data: Any, path: str) -> None:
 
 
 def torch_load(path: str, map_location: Any = None) -> Any:
+    if torch_geometric.typing.WITH_PT24:
+        try:
+            with fsspec.open(path, 'rb') as f:
+                return torch.load(f, map_location, weights_only=True)
+        except pickle.UnpicklingError as e:
+            error_msg = str(e)
+            if "add_safe_globals" in error_msg:
+                warn_msg = ("Weights only load failed. Please file an issue "
+                            "to make `torch.load(weights_only=True)` "
+                            "compatible in your case.")
+                match = re.search(r'add_safe_globals\(.*?\)', error_msg)
+                if match is not None:
+                    warnings.warn(f"{warn_msg} Please use "
+                                  f"`torch.serialization.{match.group()}` to "
+                                  f"allowlist this global.")
+                else:
+                    warnings.warn(warn_msg)
+
+                with fsspec.open(path, 'rb') as f:
+                    return torch.load(f, map_location, weights_only=False)
+            else:
+                raise e
+
     with fsspec.open(path, 'rb') as f:
         return torch.load(f, map_location)

torch_geometric/io/npz.py

@@ -1,7 +1,6 @@
 from typing import Any, Dict
 
 import numpy as np
-import scipy.sparse as sp
 import torch
 
 from torch_geometric.data import Data
@@ -15,6 +14,8 @@ def read_npz(path: str, to_undirected: bool = True) -> Data:
 
 
 def parse_npz(f: Dict[str, Any], to_undirected: bool = True) -> Data:
+    import scipy.sparse as sp
+
     x = sp.csr_matrix((f['attr_data'], f['attr_indices'], f['attr_indptr']),
                       f['attr_shape']).todense()
     x = torch.from_numpy(x).to(torch.float)

torch_geometric/io/off.py

@@ -16,7 +16,7 @@ def parse_off(src: List[str]) -> Data:
     else:
         src[0] = src[0][3:]
 
-    num_nodes, num_faces = [int(item) for item in src[0].split()[:2]]
+    num_nodes, num_faces = (int(item) for item in src[0].split()[:2])
 
     pos = parse_txt_array(src[1:1 + num_nodes])
 
@@ -52,7 +52,7 @@ def read_off(path: str) -> Data:
     Args:
         path (str): The path to the file.
     """
-    with open(path, 'r') as f:
+    with open(path) as f:
         src = f.read().split('\n')[:-1]
     return parse_off(src)
 

torch_geometric/io/sdf.py

@@ -9,7 +9,7 @@ elems = {'H': 0, 'C': 1, 'N': 2, 'O': 3, 'F': 4}
 
 def parse_sdf(src: str) -> Data:
     lines = src.split('\n')[3:]
-    num_atoms, num_bonds = [int(item) for item in lines[0].split()[:2]]
+    num_atoms, num_bonds = (int(item) for item in lines[0].split()[:2])
 
     atom_block = lines[1:num_atoms + 1]
     pos = parse_txt_array(atom_block, end=3)
@@ -28,5 +28,5 @@ def parse_sdf(src: str) -> Data:
 
 
 def read_sdf(path: str) -> Data:
-    with open(path, 'r') as f:
+    with open(path) as f:
         return parse_sdf(f.read())

torch_geometric/io/tu.py

@@ -1,7 +1,6 @@
 import os.path as osp
 from typing import Dict, List, Optional, Tuple
 
-import numpy as np
 import torch
 from torch import Tensor
 
@@ -37,7 +36,7 @@ def read_tu_data(
         if node_label.dim() == 1:
             node_label = node_label.unsqueeze(-1)
         node_label = node_label - node_label.min(dim=0)[0]
-        node_labels = node_label.unbind(dim=-1)
+        node_labels = list(node_label.unbind(dim=-1))
         node_labels = [one_hot(x) for x in node_labels]
         if len(node_labels) == 1:
             node_label = node_labels[0]
@@ -56,7 +55,7 @@ def read_tu_data(
         if edge_label.dim() == 1:
             edge_label = edge_label.unsqueeze(-1)
         edge_label = edge_label - edge_label.min(dim=0)[0]
-        edge_labels = edge_label.unbind(dim=-1)
+        edge_labels = list(edge_label.unbind(dim=-1))
         edge_labels = [one_hot(e) for e in edge_labels]
         if len(edge_labels) == 1:
             edge_label = edge_labels[0]
@@ -108,11 +107,11 @@ def cat(seq: List[Optional[Tensor]]) -> Optional[Tensor]:
 
 
 def split(data: Data, batch: Tensor) -> Tuple[Data, Dict[str, Tensor]]:
-    node_slice = cumsum(torch.from_numpy(np.bincount(batch)))
+    node_slice = cumsum(torch.bincount(batch))
 
     assert data.edge_index is not None
     row, _ = data.edge_index
-    edge_slice = cumsum(torch.from_numpy(np.bincount(batch[row])))
+    edge_slice = cumsum(torch.bincount(batch[row]))
 
     # Edge indices should start at zero for every graph.
     data.edge_index -= node_slice[batch[row]].unsqueeze(0)

torch_geometric/loader/__init__.py

@@ -22,6 +22,7 @@ from .dynamic_batch_sampler import DynamicBatchSampler
 from .prefetch import PrefetchLoader
 from .cache import CachedLoader
 from .mixin import AffinityMixin
+from .rag_loader import RAGQueryLoader, RAGFeatureStore, RAGGraphStore
 
 __all__ = classes = [
     'DataLoader',
@@ -50,6 +51,9 @@ __all__ = classes = [
     'PrefetchLoader',
     'CachedLoader',
     'AffinityMixin',
+    'RAGQueryLoader',
+    'RAGFeatureStore',
+    'RAGGraphStore'
 ]
 
 RandomNodeSampler = deprecated(

torch_geometric/loader/cluster.py

@@ -1,4 +1,5 @@
 import copy
+import os
 import os.path as osp
 import sys
 from dataclasses import dataclass
@@ -10,10 +11,11 @@ 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
-from torch_geometric.utils.sparse import index2ptr, ptr2index
 
 
 @dataclass
@@ -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/utils.py

@@ -8,7 +8,6 @@ import torch
 from torch import Tensor
 
 import torch_geometric.typing
-from torch_geometric import EdgeIndex
 from torch_geometric.data import (
     Data,
     FeatureStore,
@@ -105,13 +104,15 @@ def filter_edge_store_(store: EdgeStorage, out_store: EdgeStorage, row: Tensor,
     # which represents the new graph as denoted by `(row, col)`:
     for key, value in store.items():
         if key == 'edge_index':
+            edge_index = torch.stack([row, col], dim=0).to(value.device)
             # TODO Integrate `EdgeIndex` into `custom_store`.
-            out_store.edge_index = EdgeIndex(
-                torch.stack([row, col], dim=0).to(value.device),
-                sparse_size=out_store.size(),
-                sort_order='col',
-                # TODO Support `is_undirected`.
-            )
+            # edge_index = EdgeIndex(
+            #     torch.stack([row, col], dim=0).to(value.device),
+            #     sparse_size=out_store.size(),
+            #     sort_order='col',
+            #     # TODO Support `is_undirected`.
+            # )
+            out_store.edge_index = edge_index
 
         elif key == 'adj_t':
             # NOTE: We expect `(row, col)` to be sorted by `col` (CSC layout).

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