pyg-nightly 2.7.0.dev20250223__tar.gz → 2.7.0.dev20250225__tar.gz

{pyg_nightly-2.7.0.dev20250223 → pyg_nightly-2.7.0.dev20250225}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyg-nightly
-Version: 2.7.0.dev20250223
+Version: 2.7.0.dev20250225
 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>

{pyg_nightly-2.7.0.dev20250223 → pyg_nightly-2.7.0.dev20250225}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend="flit_core.buildapi"
 
 [project]
 name="pyg-nightly"
-version="2.7.0.dev20250223"
+version="2.7.0.dev20250225"
 authors=[
     {name="Matthias Fey", email="matthias@pyg.org"},
 ]

{pyg_nightly-2.7.0.dev20250223 → pyg_nightly-2.7.0.dev20250225}/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.dev20250223'
+__version__ = '2.7.0.dev20250225'
 
 __all__ = [
     'Index',

{pyg_nightly-2.7.0.dev20250223 → pyg_nightly-2.7.0.dev20250225}/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
@@ -245,7 +246,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,14 +300,50 @@ 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,
         index: Any,
     ) -> Union['HashTensor', Tensor]:
+        """"""  # noqa: D419
         return torch.index_select(self, dim, index)
 
     def select(  # type: ignore
@@ -314,14 +351,99 @@ class HashTensor(Tensor):
         dim: int,
         index: Any,
     ) -> Union['HashTensor', Tensor]:
+        """"""  # noqa: D419
         return torch.select(self, dim, index)
 
+    def share_memory_(self) -> 'HashTensor':
+        """"""  # noqa: D419
+        if isinstance(self._map, Tensor):
+            self._map.share_memory_()
+        if self._value is not None:
+            self._value.share_memory_()
+        self._min_key.share_memory_()
+        self._max_key.share_memory_()
+        return self
+
+    def is_shared(self) -> bool:
+        """"""  # noqa: D419
+        return self._min_key.is_shared()
+
+    def detach_(self) -> 'HashTensor':  # type: ignore
+        """"""  # noqa: D419
+        if self._value is not None:
+            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:
     return tensor._shallow_copy()
 
 
+@implements(aten.clone.default)
+def _clone(
+    tensor: HashTensor,
+    *,
+    memory_format: torch.memory_format = torch.preserve_format,
+) -> HashTensor:
+
+    value = tensor._value
+    if value is not None:
+        value = aten.clone.default(value, memory_format=memory_format)
+
+    return tensor._from_data(
+        tensor._map,  # NOTE No reason to do clone since it is read-only.
+        value,
+        tensor._min_key,  # NOTE No reason to do clone since it is read-only.
+        tensor._max_key,  # NOTE No reason to do clone since it is read-only.
+        num_keys=tensor.size(0),
+        dtype=tensor.dtype,
+    )
+
+
+@implements(aten.detach.default)
+def _detach(tensor: HashTensor) -> HashTensor:
+    value = tensor._value
+    if value is not None:
+        value = aten.detach.default(value)
+
+    return tensor._from_data(
+        tensor._map,
+        value,
+        tensor._min_key,
+        tensor._max_key,
+        num_keys=tensor.size(0),
+        dtype=tensor.dtype,
+    )
+
+
 @implements(aten._to_copy.default)
 def _to_copy(
     tensor: HashTensor,
@@ -354,8 +476,8 @@ def _to_copy(
         _map = aten._to_copy.default(_map, device=device)
     # Only convert `_map` in case `CUDAHashMap` exists - otherwise we use
     # CPU-based mapping anyway and there is no need for a copy.
-    elif (torch_geometric.typing.WITH_CUDA_HASH_MAP and min_key.is_cuda
-          and tensor._min_key.device != min_key.device):
+    elif (torch_geometric.typing.WITH_CUDA_HASH_MAP and tensor.is_cuda
+          and tensor.device != min_key.device):
         key = _map.keys()
         key = aten._to_copy.default(key, device=device)
         _map = get_hash_map(key)
@@ -370,11 +492,32 @@ def _to_copy(
     )
 
 
+@implements(aten._pin_memory.default)
+def _pin_memory(tensor: HashTensor) -> HashTensor:
+    _map = tensor._map
+    if isinstance(_map, Tensor):
+        _map = aten._pin_memory.default(_map)
+
+    value = tensor._value
+    if value is not None:
+        value = aten._pin_memory.default(value)
+
+    return tensor._from_data(
+        _map,
+        value,
+        aten._pin_memory.default(tensor._min_key),
+        aten._pin_memory.default(tensor._max_key),
+        num_keys=tensor.size(0),
+        dtype=tensor.dtype,
+    )
+
+
 @implements(aten.unsqueeze.default)
 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,
@@ -475,22 +618,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
@@ -523,20 +676,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
@@ -553,7 +711,7 @@ def _select(
         key = torch.tensor(
             [index],
             dtype=tensor._min_key.dtype,
-            device=tensor._min_key.device,
+            device=tensor.device,
         )
         return tensor._get(key).squeeze(0)
 
@@ -565,3 +723,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,
+    )