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

torch_geometric/index.py

@@ -12,6 +12,7 @@ from typing import (
     Union,
 )
 
+import numpy as np
 import torch
 import torch.utils._pytree as pytree
 from torch import Tensor
@@ -103,10 +104,10 @@ class Index(Tensor):
     conversion in case its representation is sorted.
     Caches are filled based on demand (*e.g.*, when calling
     :meth:`Index.get_indptr`), or when explicitly requested via
-    :meth:`Index.fill_cache_`, and are maintaned and adjusted over its
+    :meth:`Index.fill_cache_`, and are maintained 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 +121,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
 
@@ -181,7 +182,7 @@ class Index(Tensor):
         assert_one_dimensional(data)
         assert_contiguous(data)
 
-        out = Tensor._make_wrapper_subclass(  # type: ignore
+        out = Tensor._make_wrapper_subclass(
             cls,
             size=data.size(),
             strides=data.stride(),
@@ -360,10 +361,10 @@ class Index(Tensor):
         return index
 
     # Prevent auto-wrapping outputs back into the proper subclass type:
-    __torch_function__ = torch._C._disabled_torch_function_impl
+    __torch_function__ = torch._C._disabled_torch_function_impl  # type: ignore
 
     @classmethod
-    def __torch_dispatch__(
+    def __torch_dispatch__(  # type: ignore
         cls: Type,
         func: Callable[..., Any],
         types: Iterable[Type[Any]],
@@ -410,6 +411,14 @@ class Index(Tensor):
         return torch._tensor_str._add_suffixes(prefix + tensor_str, suffixes,
                                                indent, force_newline=False)
 
+    def tolist(self) -> List[Any]:
+        """"""  # noqa: D419
+        return self._data.tolist()
+
+    def numpy(self, *, force: bool = False) -> np.ndarray:
+        """"""  # noqa: D419
+        return self._data.numpy(force=force)
+
     # Helpers #################################################################
 
     def _shallow_copy(self) -> 'Index':
@@ -632,7 +641,7 @@ def _slice(
     step: int = 1,
 ) -> Index:
 
-    if ((start is None or start <= 0)
+    if ((start is None or start <= 0 or start <= -input.size(dim))
             and (end is None or end > input.size(dim)) and step == 1):
         return input._shallow_copy()  # No-op.
 

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
 
@@ -211,5 +214,29 @@ 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.", stacklevel=2)
+                else:
+                    warnings.warn(warn_msg, stacklevel=2)
+
+                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/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/llm/__init__.py

@@ -0,0 +1,9 @@
+from .large_graph_indexer import LargeGraphIndexer
+from .rag_loader import RAGQueryLoader
+from .utils import *  # noqa
+from .models import *  # noqa
+
+__all__ = classes = [
+    'LargeGraphIndexer',
+    'RAGQueryLoader',
+]