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

torch_geometric/datasets/word_net.py

@@ -67,7 +67,7 @@ class WordNet18(InMemoryDataset):
     def process(self) -> None:
         srcs, dsts, edge_types = [], [], []
         for path in self.raw_paths:
-            with open(path, 'r') as f:
+            with open(path) as f:
                 edges = [int(x) for x in f.read().split()[1:]]
                 edge = torch.tensor(edges, dtype=torch.long)
                 srcs.append(edge[::3])
@@ -173,7 +173,7 @@ class WordNet18RR(InMemoryDataset):
 
         srcs, dsts, edge_types = [], [], []
         for path in self.raw_paths:
-            with open(path, 'r') as f:
+            with open(path) as f:
                 edges = f.read().split()
 
                 _src = edges[::3]

torch_geometric/datasets/yelp.py

@@ -3,7 +3,6 @@ import os.path as osp
 from typing import Callable, List, Optional
 
 import numpy as np
-import scipy.sparse as sp
 import torch
 
 from torch_geometric.data import Data, InMemoryDataset, download_google_url
@@ -73,6 +72,8 @@ class Yelp(InMemoryDataset):
         download_google_url(self.role_id, self.raw_dir, 'role.json')
 
     def process(self) -> None:
+        import scipy.sparse as sp
+
         f = np.load(osp.join(self.raw_dir, 'adj_full.npz'))
         adj = sp.csr_matrix((f['data'], f['indices'], f['indptr']), f['shape'])
         adj = adj.tocoo()

torch_geometric/datasets/zinc.py

@@ -139,7 +139,7 @@ class ZINC(InMemoryDataset):
             indices = list(range(len(mols)))
 
             if self.subset:
-                with open(osp.join(self.raw_dir, f'{split}.index'), 'r') as f:
+                with open(osp.join(self.raw_dir, f'{split}.index')) as f:
                     indices = [int(x) for x in f.read()[:-1].split(',')]
 
             pbar = tqdm(total=len(indices))

torch_geometric/device.py

@@ -0,0 +1,42 @@
+from typing import Any
+
+import torch
+
+
+def is_mps_available() -> bool:
+    r"""Returns a bool indicating if MPS is currently available."""
+    if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+        try:  # Github CI may not have access to MPS hardware. Confirm:
+            torch.empty(1, device='mps')
+            return True
+        except Exception:
+            return False
+    return False
+
+
+def is_xpu_available() -> bool:
+    r"""Returns a bool indicating if XPU is currently available."""
+    if hasattr(torch, 'xpu') and torch.xpu.is_available():
+        return True
+    try:
+        import intel_extension_for_pytorch as ipex
+        return ipex.xpu.is_available()
+    except ImportError:
+        return False
+
+
+def device(device: Any) -> torch.device:
+    r"""Returns a :class:`torch.device`.
+
+    If :obj:`"auto"` is specified, returns the optimal device depending on
+    available hardware.
+    """
+    if device != 'auto':
+        return torch.device(device)
+    if torch.cuda.is_available():
+        return torch.device('cuda')
+    if is_mps_available():
+        return torch.device('mps')
+    if is_xpu_available():
+        return torch.device('xpu')
+    return torch.device('cpu')

torch_geometric/distributed/local_feature_store.py

@@ -15,6 +15,7 @@ from torch_geometric.distributed.rpc import (
     rpc_async,
     rpc_register,
 )
+from torch_geometric.io import fs
 from torch_geometric.typing import EdgeType, NodeOrEdgeType, NodeType
 
 
@@ -415,11 +416,11 @@ class LocalFeatureStore(FeatureStore):
 
         node_feats: Optional[Dict[str, Any]] = None
         if osp.exists(osp.join(part_dir, 'node_feats.pt')):
-            node_feats = torch.load(osp.join(part_dir, 'node_feats.pt'))
+            node_feats = fs.torch_load(osp.join(part_dir, 'node_feats.pt'))
 
         edge_feats: Optional[Dict[str, Any]] = None
         if osp.exists(osp.join(part_dir, 'edge_feats.pt')):
-            edge_feats = torch.load(osp.join(part_dir, 'edge_feats.pt'))
+            edge_feats = fs.torch_load(osp.join(part_dir, 'edge_feats.pt'))
 
         if not meta['is_hetero'] and node_feats is not None:
             feat_store.put_global_id(node_feats['global_id'], group_name=None)

torch_geometric/distributed/local_graph_store.py

@@ -6,6 +6,7 @@ from torch import Tensor
 
 from torch_geometric.data import EdgeAttr, GraphStore
 from torch_geometric.distributed.partition import load_partition_info
+from torch_geometric.io import fs
 from torch_geometric.typing import EdgeTensorType, EdgeType, NodeType
 from torch_geometric.utils import sort_edge_index
 
@@ -185,7 +186,7 @@ class LocalGraphStore(GraphStore):
         graph_store.edge_pb = edge_pb
         graph_store.meta = meta
 
-        graph_data = torch.load(osp.join(part_dir, 'graph.pt'))
+        graph_data = fs.torch_load(osp.join(part_dir, 'graph.pt'))
         graph_store.is_sorted = meta['is_sorted']
 
         if not meta['is_hetero']:

torch_geometric/distributed/partition.py

@@ -3,15 +3,16 @@ import logging
 import os
 import os.path as osp
 from collections import defaultdict
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 
 import torch_geometric.distributed as pyg_dist
 from torch_geometric.data import Data, HeteroData
+from torch_geometric.io import fs
 from torch_geometric.loader.cluster import ClusterData
 from torch_geometric.sampler.utils import sort_csc
-from torch_geometric.typing import Dict, EdgeType, EdgeTypeStr, NodeType, Tuple
+from torch_geometric.typing import EdgeType, EdgeTypeStr, NodeType
 
 
 class Partitioner:
@@ -23,7 +24,7 @@ class Partitioner:
 
     **Homogeneous graphs:**
 
-    .. code-block::
+    .. code-block:: none
 
         root/
         |-- META.json
@@ -40,7 +41,7 @@ class Partitioner:
 
     **Heterogeneous graphs:**
 
-    .. code-block::
+    .. code-block:: none
 
         root/
         |-- META.json
@@ -380,21 +381,21 @@ def load_partition_info(
     assert osp.exists(partition_dir)
 
     if meta['is_hetero'] is False:
-        node_pb = torch.load(osp.join(root_dir, 'node_map.pt'))
-        edge_pb = torch.load(osp.join(root_dir, 'edge_map.pt'))
+        node_pb = fs.torch_load(osp.join(root_dir, 'node_map.pt'))
+        edge_pb = fs.torch_load(osp.join(root_dir, 'edge_map.pt'))
 
         return (meta, num_partitions, partition_idx, node_pb, edge_pb)
     else:
         node_pb_dict = {}
         node_pb_dir = osp.join(root_dir, 'node_map')
         for ntype in meta['node_types']:
-            node_pb_dict[ntype] = torch.load(
+            node_pb_dict[ntype] = fs.torch_load(
                 osp.join(node_pb_dir, f'{pyg_dist.utils.as_str(ntype)}.pt'))
 
         edge_pb_dict = {}
         edge_pb_dir = osp.join(root_dir, 'edge_map')
         for etype in meta['edge_types']:
-            edge_pb_dict[tuple(etype)] = torch.load(
+            edge_pb_dict[tuple(etype)] = fs.torch_load(
                 osp.join(edge_pb_dir, f'{pyg_dist.utils.as_str(etype)}.pt'))
 
         return (meta, num_partitions, partition_idx, node_pb_dict,

torch_geometric/edge_index.py

@@ -173,7 +173,7 @@ class EdgeIndex(Tensor):
     :meth:`EdgeIndex.fill_cache_`, and are maintained and adjusted over its
     lifespan (*e.g.*, when calling :meth:`EdgeIndex.flip`).
 
-    This representation ensures for optimal computation in GNN message passing
+    This representation ensures optimal computation in GNN message passing
     schemes, while preserving the ease-of-use of regular COO-based :pyg:`PyG`
     workflows.
 
@@ -820,19 +820,28 @@ class EdgeIndex(Tensor):
                 :obj:`1.0`. (default: :obj:`None`)
         """
         value = self._get_value() if value is None else value
-        out = torch.sparse_coo_tensor(
+
+        if not torch_geometric.typing.WITH_PT21:
+            out = torch.sparse_coo_tensor(
+                indices=self._data,
+                values=value,
+                size=self.get_sparse_size(),
+                device=self.device,
+                requires_grad=value.requires_grad,
+            )
+            if self.is_sorted_by_row:
+                out = out._coalesced_(True)
+            return out
+
+        return torch.sparse_coo_tensor(
             indices=self._data,
             values=value,
             size=self.get_sparse_size(),
             device=self.device,
             requires_grad=value.requires_grad,
+            is_coalesced=True if self.is_sorted_by_row else None,
         )
 
-        if self.is_sorted_by_row:
-            out = out._coalesced_(True)
-
-        return out
-
     def to_sparse_csr(  # type: ignore
             self,
             value: Optional[Tensor] = None,
@@ -1928,7 +1937,7 @@ def _spmm(
     if transpose and not input.is_sorted_by_col:
         cls_name = input.__class__.__name__
         raise ValueError(f"'matmul(..., transpose=True)' requires "
-                         f"'{cls_name}' to be sorted by colums")
+                         f"'{cls_name}' to be sorted by columns")
 
     if (torch_geometric.typing.WITH_TORCH_SPARSE and not is_compiling()
             and other.is_cuda):  # pragma: no cover

torch_geometric/explain/algorithm/base.py

@@ -50,7 +50,6 @@ class ExplainerAlgorithm(torch.nn.Module):
         r"""Checks if the explainer supports the user-defined settings provided
         in :obj:`self.explainer_config`, :obj:`self.model_config`.
         """
-        pass
 
     ###########################################################################
 

torch_geometric/explain/algorithm/pg_explainer.py

@@ -59,7 +59,7 @@ class PGExplainer(ExplainerAlgorithm):
         'edge_size': 0.05,
         'edge_ent': 1.0,
         'temp': [5.0, 2.0],
-        'bias': 0.0,
+        'bias': 0.01,
     }
 
     def __init__(self, epochs: int, lr: float = 0.003, **kwargs):

torch_geometric/explain/explanation.py

@@ -340,10 +340,10 @@ class HeteroExplanation(HeteroData, ExplanationMixin):
         """
         node_mask_dict = self.node_mask_dict
         for node_mask in node_mask_dict.values():
-            if node_mask.dim() != 2 or node_mask.size(1) <= 1:
+            if node_mask.dim() != 2:
                 raise ValueError(f"Cannot compute feature importance for "
                                  f"object-level 'node_mask' "
-                                 f"(got shape {node_mask_dict.size()})")
+                                 f"(got shape {node_mask.size()})")
 
         if feat_labels is None:
             feat_labels = {}

torch_geometric/graphgym/checkpoint.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, List, Optional, Union
 import torch
 
 from torch_geometric.graphgym.config import cfg
+from torch_geometric.io import fs
 
 MODEL_STATE = 'model_state'
 OPTIMIZER_STATE = 'optimizer_state'
@@ -25,7 +26,7 @@ def load_ckpt(
     if not osp.exists(path):
         return 0
 
-    ckpt = torch.load(path)
+    ckpt = fs.torch_load(path)
     model.load_state_dict(ckpt[MODEL_STATE])
     if optimizer is not None and OPTIMIZER_STATE in ckpt:
         optimizer.load_state_dict(ckpt[OPTIMIZER_STATE])

torch_geometric/graphgym/logger.py

@@ -19,7 +19,7 @@ def set_printing():
     logging.root.handlers = []
     logging_cfg = {'level': logging.INFO, 'format': '%(message)s'}
     os.makedirs(cfg.run_dir, exist_ok=True)
-    h_file = logging.FileHandler('{}/logging.log'.format(cfg.run_dir))
+    h_file = logging.FileHandler(f'{cfg.run_dir}/logging.log')
     h_stdout = logging.StreamHandler(sys.stdout)
     if cfg.print == 'file':
         logging_cfg['handlers'] = [h_file]
@@ -40,7 +40,7 @@ class Logger:
         self._epoch_total = cfg.optim.max_epoch
         self._time_total = 0  # won't be reset
 
-        self.out_dir = '{}/{}'.format(cfg.run_dir, name)
+        self.out_dir = f'{cfg.run_dir}/{name}'
         os.makedirs(self.out_dir, exist_ok=True)
         if cfg.tensorboard_each_run:
             from tensorboardX import SummaryWriter
@@ -210,9 +210,9 @@ class Logger:
             }
 
         # print
-        logging.info('{}: {}'.format(self.name, stats))
+        logging.info(f'{self.name}: {stats}')
         # json
-        dict_to_json(stats, '{}/stats.json'.format(self.out_dir))
+        dict_to_json(stats, f'{self.out_dir}/stats.json')
         # tensorboard
         if cfg.tensorboard_each_run:
             dict_to_tb(stats, self.tb_writer, cur_epoch)

torch_geometric/graphgym/loss.py

@@ -10,7 +10,7 @@ def compute_loss(pred, true):
 
     Args:
         pred (torch.tensor): Unnormalized prediction
-        true (torch.tensor): Grou
+        true (torch.tensor): Ground truth labels
 
     Returns: Loss, normalized prediction score
 

torch_geometric/graphgym/utils/agg_runs.py

@@ -54,7 +54,7 @@ def agg_dict_list(dict_list):
         if key != 'epoch':
             value = np.array([dict[key] for dict in dict_list])
             dict_agg[key] = np.mean(value).round(cfg.round)
-            dict_agg['{}_std'.format(key)] = np.std(value).round(cfg.round)
+            dict_agg[f'{key}_std'] = np.std(value).round(cfg.round)
     return dict_agg
 
 
@@ -107,7 +107,7 @@ def agg_runs(dir, metric_best='auto'):
                     [stats[metric] for stats in stats_list])
                 best_epoch = \
                     stats_list[
-                        eval("performance_np.{}()".format(cfg.metric_agg))][
+                        eval(f"performance_np.{cfg.metric_agg}()")][
                         'epoch']
                 print(best_epoch)
 
@@ -190,7 +190,7 @@ def agg_batch(dir, metric_best='auto'):
             results[key] = pd.DataFrame(results[key])
             results[key] = results[key].sort_values(
                 list(dict_name.keys()), ascending=[True] * len(dict_name))
-            fname = osp.join(dir_out, '{}_best.csv'.format(key))
+            fname = osp.join(dir_out, f'{key}_best.csv')
             results[key].to_csv(fname, index=False)
 
     results = {'train': [], 'val': [], 'test': []}
@@ -213,7 +213,7 @@ def agg_batch(dir, metric_best='auto'):
             results[key] = pd.DataFrame(results[key])
             results[key] = results[key].sort_values(
                 list(dict_name.keys()), ascending=[True] * len(dict_name))
-            fname = osp.join(dir_out, '{}.csv'.format(key))
+            fname = osp.join(dir_out, f'{key}.csv')
             results[key].to_csv(fname, index=False)
 
     results = {'train': [], 'val': [], 'test': []}
@@ -245,7 +245,7 @@ def agg_batch(dir, metric_best='auto'):
             results[key] = pd.DataFrame(results[key])
             results[key] = results[key].sort_values(
                 list(dict_name.keys()), ascending=[True] * len(dict_name))
-            fname = osp.join(dir_out, '{}_bestepoch.csv'.format(key))
+            fname = osp.join(dir_out, f'{key}_bestepoch.csv')
             results[key].to_csv(fname, index=False)
 
-    print('Results aggregated across models saved in {}'.format(dir_out))
+    print(f'Results aggregated across models saved in {dir_out}')

torch_geometric/index.py

@@ -106,7 +106,7 @@ class Index(Tensor):
     :meth:`Index.fill_cache_`, and are maintaned and adjusted over its
     lifespan.
 
-    This representation ensures for optimal computation in GNN message passing
+    This representation ensures optimal computation in GNN message passing
     schemes, while preserving the ease-of-use of regular COO-based :pyg:`PyG`
     workflows.
 
@@ -120,7 +120,7 @@ class Index(Tensor):
         assert index.is_sorted
 
         # Flipping order:
-        edge_index.flip(0)
+        index.flip(0)
         >>> Index([[2, 1, 1, 0], dim_size=3)
         assert not index.is_sorted
 
@@ -685,14 +685,14 @@ def _index(
 
 @implements(aten.add.Tensor)
 def _add(
-    input: Index,
+    input: Union[int, Tensor, Index],
     other: Union[int, Tensor, Index],
     *,
     alpha: int = 1,
 ) -> Union[Index, Tensor]:
 
     data = aten.add.Tensor(
-        input._data,
+        input._data if isinstance(input, Index) else input,
         other._data if isinstance(other, Index) else other,
         alpha=alpha,
     )
@@ -704,15 +704,25 @@ def _add(
 
     out = Index(data)
 
+    if isinstance(input, Tensor) and input.numel() <= 1:
+        input = int(input)
+
     if isinstance(other, Tensor) and other.numel() <= 1:
         other = int(other)
 
     if isinstance(other, int):
+        assert isinstance(input, Index)
         if input.dim_size is not None:
             out._dim_size = input.dim_size + alpha * other
         out._is_sorted = input.is_sorted
 
-    elif isinstance(other, Index):
+    elif isinstance(input, int):
+        assert isinstance(other, Index)
+        if other.dim_size is not None:
+            out._dim_size = input + alpha * other.dim_size
+        out._is_sorted = other.is_sorted
+
+    elif isinstance(input, Index) and isinstance(other, Index):
         if input.dim_size is not None and other.dim_size is not None:
             out._dim_size = input.dim_size + alpha * other.dim_size
 
@@ -754,14 +764,14 @@ def add_(
 
 @implements(aten.sub.Tensor)
 def _sub(
-    input: Index,
+    input: Union[int, Tensor, Index],
     other: Union[int, Tensor, Index],
     *,
     alpha: int = 1,
 ) -> Union[Index, Tensor]:
 
     data = aten.sub.Tensor(
-        input._data,
+        input._data if isinstance(input, Index) else input,
         other._data if isinstance(other, Index) else other,
         alpha=alpha,
     )
@@ -773,6 +783,9 @@ def _sub(
 
     out = Index(data)
 
+    if not isinstance(input, Index):
+        return out
+
     if isinstance(other, Tensor) and other.numel() <= 1:
         other = int(other)
 

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`)
         """
@@ -448,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
 
@@ -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(