pyg-nightly 2.7.0.dev20250414__py3-none-any.whl → 2.7.0.dev20250416__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250414.dist-info → pyg_nightly-2.7.0.dev20250416.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyg-nightly
-Version: 2.7.0.dev20250414
+Version: 2.7.0.dev20250416
 Summary: Graph Neural Network Library for PyTorch
 Keywords: deep-learning,pytorch,geometric-deep-learning,graph-neural-networks,graph-convolutional-networks
 Author-email: Matthias Fey <matthias@pyg.org>
@@ -61,6 +61,12 @@ Requires-Dist: protobuf<4.21 ; extra == "graphgym"
 Requires-Dist: pytorch-lightning<2.3.0 ; extra == "graphgym"
 Requires-Dist: yacs ; extra == "graphgym"
 Requires-Dist: huggingface_hub ; extra == "modelhub"
+Requires-Dist: pcst_fast ; extra == "rag"
+Requires-Dist: datasets ; extra == "rag"
+Requires-Dist: transformers ; extra == "rag"
+Requires-Dist: pandas ; extra == "rag"
+Requires-Dist: sentencepiece ; extra == "rag"
+Requires-Dist: accelerate ; extra == "rag"
 Requires-Dist: onnx ; extra == "test"
 Requires-Dist: onnxruntime ; extra == "test"
 Requires-Dist: pytest ; extra == "test"
@@ -74,6 +80,7 @@ Provides-Extra: dev
 Provides-Extra: full
 Provides-Extra: graphgym
 Provides-Extra: modelhub
+Provides-Extra: rag
 Provides-Extra: test
 
 <p align="center">

{pyg_nightly-2.7.0.dev20250414.dist-info → pyg_nightly-2.7.0.dev20250416.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=gL3VBzLbnTDZMx7XwpSf7shcnqbMg36uT_MEgesPhSY,1978
+torch_geometric/__init__.py,sha256=en3YAn8JWtZUWaMs1gx9X-Z0xZM9IfGB3fvTW53cJno,1978
 torch_geometric/_compile.py,sha256=f-WQeH4VLi5Hn9lrgztFUCSrN_FImjhQa6BxFzcYC38,1338
 torch_geometric/_onnx.py,sha256=V9ffrIKSqhDw6xUZ12lkuSfNs48cQp2EeJ6Z19GfnVw,349
 torch_geometric/backend.py,sha256=lVaf7aLoVaB3M-UcByUJ1G4T4FOK6LXAg0CF4W3E8jo,1575
@@ -44,7 +44,7 @@ torch_geometric/data/graph_store.py,sha256=EtIgsyY7RdBHRTCn34VypEBOG8cg8WzsNT_kT
 torch_geometric/data/hetero_data.py,sha256=q0L3bENyEvo_BGLPwZPVzh730Aak6sQ7yXoawPgM72E,47982
 torch_geometric/data/hypergraph_data.py,sha256=33hsXW25Yz4Ju8mKajYinZOrkqrUi1SqThG7MlOOYNM,8294
 torch_geometric/data/in_memory_dataset.py,sha256=F35hU9Dw3qiJUL5E1CCAfq-1xrlUMstXBmQVEQdtJ1I,13403
-torch_geometric/data/large_graph_indexer.py,sha256=eor7F98kPDlrs1v0qypinVbhQr8wnw8mAfsLNiNmEwY,25390
+torch_geometric/data/large_graph_indexer.py,sha256=3mF2c7BLbnAg8h5mBaG95_lyQMiXJ_uhEe5bXbX7e1M,25430
 torch_geometric/data/makedirs.py,sha256=6uOv4y34i947cm4rv7Aj2_YZBq-EOsyPKnlGA188YSw,463
 torch_geometric/data/on_disk_dataset.py,sha256=77om-e6kzcpBb77kf7um1xY8-yHmQaao_6R7I-3NwHk,6629
 torch_geometric/data/remote_backend_utils.py,sha256=Rzpq1PczXuHhUscrFtIAL6dua6pMehSJlXG7yEsrrrg,4503
@@ -441,7 +441,7 @@ torch_geometric/nn/models/g_retriever.py,sha256=CdSOasnPiMvq5AjduNTpz-LIZiNp3X0x
 torch_geometric/nn/models/git_mol.py,sha256=Wc6Hx6RDDR7sDWRWHfA5eK9e9gFsrTZ9OLmpMfoj3pE,12676
 torch_geometric/nn/models/glem.py,sha256=sT0XM4klVlci9wduvUoXupATUw9p25uXtaJBrmv3yvs,16431
 torch_geometric/nn/models/gnnff.py,sha256=15dkiLgy0LmH1hnUrpeoHioIp4BPTfjpVATpnGRt9E0,7860
-torch_geometric/nn/models/gpse.py,sha256=Fwldw9N3axV--BcSnI4im1sy1r87a5mAXZAXHu_6k2Y,41932
+torch_geometric/nn/models/gpse.py,sha256=NSN7yL98a9KzPjhwEpqOyx4RVtEfOPhFc0opejpTGVk,41981
 torch_geometric/nn/models/graph_mixer.py,sha256=mthMeCOikR8gseEsu4oJ3Cd9C35zHSv1p32ROwnG-6s,9246
 torch_geometric/nn/models/graph_unet.py,sha256=N8TSmJo8AlbZjjcame0xW_jZvMOirL5ahw6qv5Yjpbs,5586
 torch_geometric/nn/models/jumping_knowledge.py,sha256=9JR2EoViXKjcDSLb8tjJm-UHfv1mQCJvZAAEsYa0Ocw,5496
@@ -515,15 +515,15 @@ torch_geometric/sampler/base.py,sha256=kT6hYM6losYta3pqLQlqiqboJiujLy6RlH8qM--U_
 torch_geometric/sampler/hgt_sampler.py,sha256=UAm8_wwzEcziKDJ8-TnfZh1705dXRsy_I5PKhZSDTK8,2721
 torch_geometric/sampler/neighbor_sampler.py,sha256=MAVphWqNf0-cwlHRvdiU8de86dBxwjm3Miam_6s1ep4,33971
 torch_geometric/sampler/utils.py,sha256=RJtasO6Q7Pp3oYEOWrbf2DEYuSfuKZOsF2I7-eJDnoA,5485
-torch_geometric/testing/__init__.py,sha256=0mAGVWRrTBNsGV2YUkCu_FkyQ8JIcrYVw2LsdKgY9ak,1291
+torch_geometric/testing/__init__.py,sha256=m3yp_5UnCAxVgzTFofpiVt0vdbl5GwVAve8WTrAaNxo,1319
 torch_geometric/testing/asserts.py,sha256=DLC9HnBgFWuTIiQs2OalsQcXGhOVG-e6R99IWhkO32c,4606
 torch_geometric/testing/data.py,sha256=O1qo8FyNxt6RGf63Ys3eXBfa5RvYydeZLk74szrez3c,2604
-torch_geometric/testing/decorators.py,sha256=BaaQRGgU1r1_p3OTpSLzgUaROq6EB2vojq4VKSzXE6w,8657
+torch_geometric/testing/decorators.py,sha256=k7QtPWz-8Ft89cvKmZt8y1h00j2BB1yhEHwmGzTcol0,9075
 torch_geometric/testing/distributed.py,sha256=ZZCCXqiQC4-m1ExSjDZhS_a1qPXnHEwhJGTmACxNnVI,2227
 torch_geometric/testing/feature_store.py,sha256=J6JBIt2XK-t8yG8B4JzXp-aJcVl5jaCS1m2H7d6OUxs,2158
 torch_geometric/testing/graph_store.py,sha256=00B7QToCIspYmgN7svQKp1iU-qAzEtrt3VQRFxkHfuk,1044
 torch_geometric/transforms/__init__.py,sha256=P0R2CFg9pXxjTX4NnYfNPrifRPAw5lVXEOxO80q-1Ek,4296
-torch_geometric/transforms/add_gpse.py,sha256=4o0UrSmTu3CKsL3UAREiul8O4lC02PUx_ajxP4sPsxU,1570
+torch_geometric/transforms/add_gpse.py,sha256=ex7cSLmcSIn-dC5gbic-CLwKegV5Is7Y8mUn-mSOWXg,1555
 torch_geometric/transforms/add_metapaths.py,sha256=GabaPRvUnpFrZJsxLMUBY2Egzx94GTgsMxegL_qTtbk,14239
 torch_geometric/transforms/add_positional_encoding.py,sha256=tuilyubAn3yeyz8mvFc5zxXTlNzh8okKzG9AE2lPG1Q,6049
 torch_geometric/transforms/add_remaining_self_loops.py,sha256=ItU5FAcE-mkbp_wqTLkRhv0RShR5JVr8vr9d5xv3_Ak,2085
@@ -636,7 +636,7 @@ torch_geometric/utils/undirected.py,sha256=H_nfpI0_WluOG6VfjPyldvcjL4w5USAKWu2x5
 torch_geometric/visualization/__init__.py,sha256=PyR_4K5SafsJrBr6qWrkjKr6GBL1b7FtZybyXCDEVwY,154
 torch_geometric/visualization/graph.py,sha256=ZuLPL92yGRi7lxlqsUPwL_EVVXF7P2kMcveTtW79vpA,4784
 torch_geometric/visualization/influence.py,sha256=CWMvuNA_Nf1sfbJmQgn58yS4OFpeKXeZPe7kEuvkUBw,477
-pyg_nightly-2.7.0.dev20250414.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250414.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-pyg_nightly-2.7.0.dev20250414.dist-info/METADATA,sha256=1EHha-manNSei-aDwCeZuM7b0CJ05SQ4PN3dMRTPWtY,62703
-pyg_nightly-2.7.0.dev20250414.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250416.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250416.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+pyg_nightly-2.7.0.dev20250416.dist-info/METADATA,sha256=yDodf56EgttruZas0nqEgbfnFaNHm03BHQsFi-IkPf0,62979
+pyg_nightly-2.7.0.dev20250416.dist-info/RECORD,,

torch_geometric/__init__.py

@@ -31,7 +31,7 @@ from .lazy_loader import LazyLoader
 contrib = LazyLoader('contrib', globals(), 'torch_geometric.contrib')
 graphgym = LazyLoader('graphgym', globals(), 'torch_geometric.graphgym')
 
-__version__ = '2.7.0.dev20250414'
+__version__ = '2.7.0.dev20250416'
 
 __all__ = [
     'Index',

torch_geometric/data/large_graph_indexer.py

@@ -22,6 +22,7 @@ from torch import Tensor
 from tqdm import tqdm
 
 from torch_geometric.data import Data
+from torch_geometric.io import fs
 from torch_geometric.typing import WITH_PT24
 
 # Could be any hashable type
@@ -505,13 +506,13 @@ class LargeGraphIndexer:
         for fname in os.listdir(node_attr_path):
             full_fname = f"{node_attr_path}/{fname}"
             key = fname.split(".")[0]
-            indexer.node_attr[key] = torch.load(full_fname)
+            indexer.node_attr[key] = fs.torch_load(full_fname)
 
         edge_attr_path = path + "/edge_attr"
         for fname in os.listdir(edge_attr_path):
             full_fname = f"{edge_attr_path}/{fname}"
             key = fname.split(".")[0]
-            indexer.edge_attr[key] = torch.load(full_fname)
+            indexer.edge_attr[key] = fs.torch_load(full_fname)
 
         return indexer
 

torch_geometric/nn/models/gpse.py

@@ -9,6 +9,7 @@ import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch.nn import Module
 from tqdm import trange
 
 import torch_geometric.transforms as T
@@ -715,8 +716,9 @@ class GPSENodeEncoder(torch.nn.Module):
 
 
 @torch.no_grad()
-def gpse_process(model: GPSE, data: Data, rand_type: str, use_vn: bool = True,
-                 bernoulli_thresh: float = 0.5, neighbor_loader: bool = False,
+def gpse_process(model: Module, data: Data, rand_type: str,
+                 use_vn: bool = True, bernoulli_thresh: float = 0.5,
+                 neighbor_loader: bool = False,
                  num_neighbors: List[int] = [30, 20, 10], fillval: int = 5,
                  layers_mp: int = None, **kwargs) -> torch.Tensor:
     r"""Processes the data using the :class:`GPSE` model to generate and append
@@ -731,7 +733,7 @@ def gpse_process(model: GPSE, data: Data, rand_type: str, use_vn: bool = True,
     :obj:`precompute_GPSE` on your whole dataset is advised instead.
 
     Args:
-        model (GPSE): The :class:`GPSE` model.
+        model (Module): The :class:`GPSE` model.
         data (torch_geometric.data.Data): A :class:`~torch_geometric.data.Data`
             object.
         rand_type (str, optional): Type of random features to use. Options are

torch_geometric/testing/__init__.py

@@ -17,6 +17,7 @@ from .decorators import (
     onlyOnline,
     onlyGraphviz,
     onlyNeighborSampler,
+    onlyRAG,
     has_package,
     withPackage,
     withDevice,
@@ -49,6 +50,7 @@ __all__ = [
     'onlyOnline',
     'onlyGraphviz',
     'onlyNeighborSampler',
+    'onlyRAG',
     'has_package',
     'withPackage',
     'withDevice',

torch_geometric/testing/decorators.py

@@ -15,6 +15,11 @@ from torch_geometric.typing import WITH_METIS, WITH_PYG_LIB, WITH_TORCH_SPARSE
 from torch_geometric.visualization.graph import has_graphviz
 
 
+def is_rag_test() -> bool:
+    r"""Whether to run the RAG test suite."""
+    return os.getenv('RAG_TEST', '0') == '1'
+
+
 def is_full_test() -> bool:
     r"""Whether to run the full but time-consuming test suite."""
     return os.getenv('FULL_TEST', '0') == '1'
@@ -204,6 +209,18 @@ def withPackage(*args: str) -> Callable:
     return decorator
 
 
+def onlyRAG(func: Callable) -> Callable:
+    r"""A decorator to specify that this function belongs to the RAG test
+    suite.
+    """
+    import pytest
+    func = pytest.mark.rag(func)
+    return pytest.mark.skipif(
+        not is_rag_test(),
+        reason="RAG tests are disabled",
+    )(func)
+
+
 def withCUDA(func: Callable) -> Callable:
     r"""A decorator to test both on CPU and CUDA (if available)."""
     import pytest

torch_geometric/transforms/add_gpse.py

@@ -1,6 +1,7 @@
+from torch.nn import Module
+
 from torch_geometric.data import Data
 from torch_geometric.data.datapipes import functional_transform
-from torch_geometric.nn.models.gpse import GPSE
 from torch_geometric.transforms import BaseTransform, VirtualNode
 
 
@@ -13,7 +14,7 @@ class AddGPSE(BaseTransform):
     the actual encodings.
 
     Args:
-        model (GPSE): The pre-trained GPSE model.
+        model (Module): The pre-trained GPSE model.
         use_vn (bool, optional): Whether to use virtual nodes.
             (default: :obj:`True`)
         rand_type (str, optional): Type of random features to use. Options are
@@ -21,7 +22,7 @@ class AddGPSE(BaseTransform):
             (default: :obj:`NormalSE`)
 
     """
-    def __init__(self, model: GPSE, use_vn: bool = True,
+    def __init__(self, model: Module, use_vn: bool = True,
                  rand_type: str = 'NormalSE'):
         self.model = model
         self.use_vn = use_vn