pyg-nightly 2.7.0.dev20250224__py3-none-any.whl → 2.7.0.dev20250226__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250224.dist-info → pyg_nightly-2.7.0.dev20250226.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyg-nightly
-Version: 2.7.0.dev20250224
+Version: 2.7.0.dev20250226
 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>
@@ -36,7 +36,6 @@ Requires-Dist: torch_geometric[test] ; extra == "dev"
 Requires-Dist: scipy ; extra == "full"
 Requires-Dist: scikit-learn ; extra == "full"
 Requires-Dist: ase ; extra == "full"
-Requires-Dist: captum<0.7.0 ; extra == "full"
 Requires-Dist: graphviz ; extra == "full"
 Requires-Dist: h5py ; extra == "full"
 Requires-Dist: matplotlib ; extra == "full"

{pyg_nightly-2.7.0.dev20250224.dist-info → pyg_nightly-2.7.0.dev20250226.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=nc0R-kBpZZJ_9mHk-xQJhyukW9OMLPN9EmnGrJIwu1c,1978
+torch_geometric/__init__.py,sha256=MAAgbT2XVvKADBsLycGtgqDaaCWjhrtphXMHu1uUO9k,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
@@ -7,11 +7,11 @@ torch_geometric/config_store.py,sha256=zdMzlgBpUmBkPovpYQh5fMNwTZLDq2OneqX47QEx7
 torch_geometric/debug.py,sha256=cLyH9OaL2v7POyW-80b19w-ctA7a_5EZsS4aUF1wc2U,1295
 torch_geometric/deprecation.py,sha256=dWRymDIUkUVI2MeEmBG5WF4R6jObZeseSBV9G6FNfjc,858
 torch_geometric/device.py,sha256=tU5-_lBNVbVHl_kUmWPwiG5mQ1pyapwMF4JkmtNN3MM,1224
-torch_geometric/edge_index.py,sha256=BsLh5tOZRjjSYDkjqOFAdBuvMaDg7EWaaLELYsUL0Z8,70048
+torch_geometric/edge_index.py,sha256=J8uHIUXleOv4A2XONkT6yaTJi9yXauxBRJiHmJHMox4,70330
 torch_geometric/experimental.py,sha256=JbtNNEXjFGI8hZ9raM6-qrZURP6Z5nlDK8QicZUIbz0,4756
-torch_geometric/hash_tensor.py,sha256=koofBrEEo5oHCCt0gH6gCkiywyxSFSGdk-keg7MXGf4,19490
+torch_geometric/hash_tensor.py,sha256=WB-aBCJWNWqnlnzQ8Ob4LHeCXm0u1_NPPhmNAEwBpq4,24906
 torch_geometric/home.py,sha256=EV54B4Dmiv61GDbkCwtCfWGWJ4eFGwZ8s3KOgGjwYgY,790
-torch_geometric/index.py,sha256=9ChzWFCwj2slNcVBOgfV-wQn-KscJe_y7502w-Vf76w,24045
+torch_geometric/index.py,sha256=VZGVSb19biQ-HyvZA6esAPW_23dgmtQLSc3436WjP64,24327
 torch_geometric/inspector.py,sha256=nKi5o4Mn6xsG0Ex1GudTEQt_EqnF9mcMqGtp7Shh9sQ,19336
 torch_geometric/isinstance.py,sha256=truZjdU9PxSvjJ6k0d_CLJ2iOpen2o8U-54pbUbNRyE,935
 torch_geometric/lazy_loader.py,sha256=SM0UcXtIdiFge75MKBAWXedoiSOdFDOV0rm1PfoF9cE,908
@@ -633,7 +633,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.dev20250224.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250224.dist-info/WHEEL,sha256=_2ozNFCLWc93bK4WKHCO-eDUENDlo-dgc9cU3qokYO4,82
-pyg_nightly-2.7.0.dev20250224.dist-info/METADATA,sha256=OzM9N2hJ-COuEqz7HmEj8_wZkdLZDavaZfWDS5j08SA,63021
-pyg_nightly-2.7.0.dev20250224.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250226.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250226.dist-info/WHEEL,sha256=_2ozNFCLWc93bK4WKHCO-eDUENDlo-dgc9cU3qokYO4,82
+pyg_nightly-2.7.0.dev20250226.dist-info/METADATA,sha256=ohLoC35Q3PN8CDwaHoNilg0ALbm8R_VEFqF2OPyqKj4,62975
+pyg_nightly-2.7.0.dev20250226.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.dev20250224'
+__version__ = '2.7.0.dev20250226'
 
 __all__ = [
     'Index',

torch_geometric/edge_index.py

@@ -17,6 +17,7 @@ from typing import (
     overload,
 )
 
+import numpy as np
 import torch
 import torch.utils._pytree as pytree
 from torch import Tensor
@@ -1246,6 +1247,14 @@ class EdgeIndex(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) -> 'EdgeIndex':
@@ -1478,7 +1487,7 @@ def _slice(
     step: int = 1,
 ) -> Union[EdgeIndex, Tensor]:
 
-    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/hash_tensor.py

@@ -12,6 +12,7 @@ from typing import (
     Union,
 )
 
+import numpy as np
 import torch
 import torch.utils._pytree as pytree
 import xxhash
@@ -44,13 +45,6 @@ def as_key_tensor(
         key = torch.as_tensor(key, device=device)
     except Exception:
         device = device or torch.get_default_device()
-        # TODO Convert int64 to int32.
-        # On GPU, we default to int32 for faster 'CUDAHashMap' implementation:
-        if torch_geometric.typing.WITH_CUDA_HASH_MAP and device.type == 'cuda':
-            pass
-            # key = torch.tensor(
-            #     [xxhash.xxh32(x).intdigest() & 0x7FFFFFFF for x in key],
-            #     dtype=torch.int32, device=device)
         key = torch.tensor(
             [xxhash.xxh64(x).intdigest() & 0x7FFFFFFFFFFFFFFF for x in key],
             dtype=torch.int64, device=device)
@@ -72,7 +66,6 @@ def as_key_tensor(
 
 def get_hash_map(key: Tensor) -> Union[CPUHashMap, CUDAHashMap]:
     if torch_geometric.typing.WITH_CUDA_HASH_MAP and key.is_cuda:
-        # TODO Convert int64 to int32.
         return CUDAHashMap(key, 0.5)
 
     if key.is_cuda:
@@ -93,6 +86,62 @@ def get_hash_map(key: Tensor) -> Union[CPUHashMap, CUDAHashMap]:
 
 
 class HashTensor(Tensor):
+    r"""A :pytorch:`null` :class:`torch.Tensor` that can be referenced by
+    arbitrary keys rather than indices in the first dimension.
+
+    :class:`HashTensor` sub-classes a general :pytorch:`null`
+    :class:`torch.Tensor`, and extends it by CPU- and GPU-accelerated mapping
+    routines. This allow for fast and efficient access to non-contiguous
+    indices/keys while the underlying data is stored in a compact format.
+
+    This representation is ideal for scenarios where one needs a fast mapping
+    routine without relying on CPU-based external packages, and can be used,
+    *e.g.*, to perform mapping of global indices to local indices during
+    subgraph creation, or in data-processing pipelines to map non-contiguous
+    input data into a contiguous space, such as
+
+    * mapping of hashed node IDs to range :obj:`[0, num_nodes - 1]`
+    * mapping of raw input data, *e.g.*, categorical data to range
+      :obj:`[0, num_categories - 1]`
+
+    Specifically, :class:`HashTensor` supports *any* keys of *any* type,
+    *e.g.*, strings, timestamps, etc.
+
+    .. code-block:: python
+
+        from torch_geometric import HashTensor
+
+        key = torch.tensor([1000, 100, 10000])
+        value = torch.randn(3, 4)
+
+        tensor = HashTensor(key, value)
+        assert tensor.size() == (3, 4)
+
+        # Filtering:
+        query = torch.tensor([10000, 1000])
+        out = tensor[query]
+        assert out.equal(value[[2, 0]])
+
+        # Accessing non-existing keys:
+        out = tensor[[10000, 0]]
+        out.isnan()
+        >>> tensor([[False, False, False, False],
+        ...         [True, True, True, True])
+
+        # If `value` is not given, indexing returns the position of `query` in
+        # `key`, and `-1` otherwise:
+        key = ['Animation', 'Comedy', 'Fantasy']
+        tensor = HashTensor(key)
+
+        out = tensor[['Comedy', 'Romance']]
+        >>> tensor([1, -1])
+
+    Args:
+        key: The keys in the first dimension.
+        value: The values to hold.
+        dtype: The desired data type of the values of the returned tensor.
+        device: The device of the returned tensor.
+    """
     _map: Union[Tensor, CPUHashMap, CUDAHashMap]
     _value: Optional[Tensor]
     _min_key: Tensor
@@ -245,7 +294,7 @@ class HashTensor(Tensor):
             import pandas as pd
 
             ser = pd.Series(query.cpu().numpy(), dtype=self._map)
-            index = torch.from_numpy(ser.cat.codes.to_numpy()).to(torch.long)
+            index = torch.from_numpy(ser.cat.codes.to_numpy().copy()).long()
 
         index = index.to(self.device)
 
@@ -299,10 +348,44 @@ class HashTensor(Tensor):
                                           kwargs)
         return func(*args, **(kwargs or {}))
 
+    def __tensor_flatten__(self) -> Tuple[List[str], Tuple[Any, ...]]:
+        attrs = ['_map', '_min_key', '_max_key']
+        if self._value is not None:
+            attrs.append('_value')
+
+        ctx = (self.size(0), self.dtype)
+
+        return attrs, ctx
+
+    @staticmethod
+    def __tensor_unflatten__(
+        inner_tensors: Dict[str, Any],
+        ctx: Tuple[Any, ...],
+        outer_size: Tuple[int, ...],
+        outer_stride: Tuple[int, ...],
+    ) -> 'HashTensor':
+        return HashTensor._from_data(
+            inner_tensors['_map'],
+            inner_tensors.get('_value', None),
+            inner_tensors['_min_key'],
+            inner_tensors['_min_key'],
+            num_keys=ctx[0],
+            dtype=ctx[1],
+        )
+
+    def __repr__(self) -> str:  # type: ignore
+        indent = len(f'{self.__class__.__name__}(')
+        tensor_str = torch._tensor_str._tensor_str(self.as_tensor(), indent)
+        return torch._tensor_str._str_intern(self, tensor_contents=tensor_str)
+
     def tolist(self) -> List[Any]:
         """"""  # noqa: D419
         return self.as_tensor().tolist()
 
+    def numpy(self, *, force: bool = False) -> np.ndarray:
+        """"""  # noqa: D419
+        return self.as_tensor().numpy(force=force)
+
     def index_select(  # type: ignore
         self,
         dim: int,
@@ -339,6 +422,33 @@ class HashTensor(Tensor):
             self._value.detach_()
         return super().detach_()  # type: ignore
 
+    def __getitem__(self, indices: Any) -> Union['HashTensor', Tensor]:
+        if not isinstance(indices, tuple):
+            indices = (indices, )
+        assert len(indices) > 0
+
+        # We convert any index tensor in the first dimension into a tensor.
+        # This means that downstream handling (i.e. in `aten.index.Tensor`)
+        # needs to take this pre-conversion into account. However, detecting
+        # whether the first dimension is indexed can be tricky at times:
+        # * We need to take into account `Ellipsis`
+        # * We need to take any unsqueezing into account
+        if indices[0] is Ellipsis and len(indices) > 1:
+            nonempty_indices = [i for i in indices[1:] if i is not None]
+            if len(nonempty_indices) == self.dim():
+                indices = indices[1:]
+
+        if isinstance(indices[0], (int, bool)):
+            index: Union[int, Tensor] = int(as_key_tensor([indices[0]]))
+            indices = (index, ) + indices[1:]
+        elif isinstance(indices[0], (Tensor, list, np.ndarray)):
+            index = as_key_tensor(indices[0], device=self.device)
+            indices = (index, ) + indices[1:]
+
+        indices = indices[0] if len(indices) == 1 else indices
+
+        return super().__getitem__(indices)
+
 
 @implements(aten.alias.default)
 def _alias(tensor: HashTensor) -> HashTensor:
@@ -454,7 +564,8 @@ def _pin_memory(tensor: HashTensor) -> HashTensor:
 def _unsqueeze(tensor: HashTensor, dim: int) -> HashTensor:
     if dim == 0 or dim == -(tensor.dim() + 1):
         raise IndexError(f"Cannot unsqueeze '{tensor.__class__.__name__}' in "
-                         f"the first dimension")
+                         f"the first dimension. Please call `as_tensor()` "
+                         f"beforehand")
 
     return tensor._from_data(
         tensor._map,
@@ -528,7 +639,7 @@ def _slice(
 ) -> HashTensor:
 
     if dim == 0 or dim == -tensor.dim():
-        copy = start is None or (start == 0 or start <= -tensor.size(0))
+        copy = start is None or start == 0 or start <= -tensor.size(0)
         copy &= end is None or end > tensor.size(0)
         copy &= step == 1
         if copy:
@@ -555,22 +666,32 @@ _old_index_select = torch.index_select
 
 def _new_index_select(
     input: Tensor,
-    dim: int,
+    dim: Union[int, str],
     index: Tensor,
-    *,
     out: Optional[Tensor] = None,
 ) -> Tensor:
 
-    if dim < -input.dim() or dim >= input.dim():
+    if isinstance(dim, int) and (dim < -input.dim() or dim >= input.dim()):
         raise IndexError(f"Dimension out of range (expected to be in range of "
                          f"[{-input.dim()}, {input.dim()-1}], but got {dim})")
 
     # We convert any index tensor in the first dimension into a tensor. This
     # means that downstream handling (i.e. in `aten.index_select.default`)
     # needs to take this pre-conversion into account.
-    if isinstance(input, HashTensor) and (dim == 0 or dim == -input.dim()):
+    if (not torch.jit.is_scripting() and isinstance(input, HashTensor)
+            and isinstance(dim, int) and (dim == 0 or dim == -input.dim())):
         index = as_key_tensor(index, device=input.device)
-    return _old_index_select(input, dim, index, out=out)
+
+    if isinstance(dim, int):  # Type narrowing...
+        if out is None:
+            return _old_index_select(input, dim, index)
+        else:
+            return _old_index_select(input, dim, index, out=out)
+    else:
+        if out is None:
+            return _old_index_select(input, dim, index)
+        else:
+            return _old_index_select(input, dim, index, out=out)
 
 
 torch.index_select = _new_index_select  # type: ignore
@@ -603,20 +724,25 @@ _old_select = torch.select
 
 def _new_select(
     input: Tensor,
-    dim: int,
+    dim: Union[int, str],
     index: int,
 ) -> Tensor:
 
-    if dim < -input.dim() or dim >= input.dim():
+    if isinstance(dim, int) and (dim < -input.dim() or dim >= input.dim()):
         raise IndexError(f"Dimension out of range (expected to be in range of "
                          f"[{-input.dim()}, {input.dim()-1}], but got {dim})")
 
     # We convert any index in the first dimension into an integer. This means
     # that downstream handling (i.e. in `aten.select.int`) needs to take this
     # pre-conversion into account.
-    if isinstance(input, HashTensor) and (dim == 0 or dim == -input.dim()):
+    if (not torch.jit.is_scripting() and isinstance(input, HashTensor)
+            and isinstance(dim, int) and (dim == 0 or dim == -input.dim())):
         index = int(as_key_tensor([index]))
-    return _old_select(input, dim, index)
+
+    if isinstance(dim, int):  # Type narrowing...
+        return _old_select(input, dim, index)
+    else:
+        return _old_select(input, dim, index)
 
 
 torch.select = _new_select  # type: ignore
@@ -645,3 +771,27 @@ def _select(
         num_keys=tensor.size(0),
         dtype=tensor.dtype,
     )
+
+
+@implements(aten.index.Tensor)
+def _index(
+    tensor: HashTensor,
+    indices: List[Optional[Tensor]],
+) -> Union[HashTensor, Tensor]:
+
+    assert len(indices) > 0
+
+    if indices[0] is not None:
+        out = tensor._get(indices[0])
+        if len(indices) > 1:
+            out = aten.index.Tensor(out, [None] + indices[1:])
+        return out
+
+    return tensor._from_data(
+        tensor._map,
+        aten.index.Tensor(tensor.as_tensor(), indices),
+        tensor._min_key,
+        tensor._max_key,
+        num_keys=tensor.size(0),
+        dtype=tensor.dtype,
+    )

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
@@ -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.