pyg-nightly 2.7.0.dev20250222__tar.gz → 2.7.0.dev20250224__tar.gz

{pyg_nightly-2.7.0.dev20250222 → pyg_nightly-2.7.0.dev20250224}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyg-nightly
-Version: 2.7.0.dev20250222
+Version: 2.7.0.dev20250224
 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.dev20250222 → pyg_nightly-2.7.0.dev20250224}/pyproject.toml

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

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

{pyg_nightly-2.7.0.dev20250222 → pyg_nightly-2.7.0.dev20250224}/torch_geometric/hash_tensor.py

@@ -209,6 +209,20 @@ class HashTensor(Tensor):
 
         return out
 
+    @property
+    def _key(self) -> Tensor:
+        if isinstance(self._map, Tensor):
+            mask = self._map >= 0
+            key = mask.nonzero().view(-1) - 1
+            key = key[self._map[mask]]
+        elif (torch_geometric.typing.WITH_CUDA_HASH_MAP
+              or torch_geometric.typing.WITH_CPU_HASH_MAP):
+            key = self._map.keys().to(self.device)
+        else:
+            key = torch.from_numpy(self._map.categories.to_numpy())
+
+        return key.to(self.device)
+
     def _shallow_copy(self) -> 'HashTensor':
         return self._from_data(
             self._map,
@@ -239,12 +253,16 @@ class HashTensor(Tensor):
             return index.to(self.dtype)
 
         out = self._value[index]
+
         mask = index != -1
         mask = mask.view([-1] + [1] * (out.dim() - 1))
-        if out.is_floating_point():
-            return out.where(mask, float('NaN'))
+        fill_value = float('NaN') if out.is_floating_point() else -1
+        if torch_geometric.typing.WITH_PT20:
+            other: Union[int, float, Tensor] = fill_value
         else:
-            return out.where(mask, -1)
+            other = torch.full_like(out, fill_value)
+
+        return out.where(mask, other)
 
     # Methods #################################################################
 
@@ -281,15 +299,89 @@ class HashTensor(Tensor):
                                           kwargs)
         return func(*args, **(kwargs or {}))
 
-    def index_select(self, dim: int, index: Any) -> Tensor:  # type: ignore
+    def tolist(self) -> List[Any]:
+        """"""  # noqa: D419
+        return self.as_tensor().tolist()
+
+    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
+        self,
+        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
+
 
 @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,
@@ -322,8 +414,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)
@@ -338,6 +430,26 @@ 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):
@@ -359,9 +471,13 @@ def _squeeze_default(tensor: HashTensor) -> HashTensor:
     if tensor._value is None:
         return tensor._shallow_copy()
 
+    value = tensor.as_tensor()
+    for d in range(tensor.dim() - 1, 0, -1):
+        value = value.squeeze(d)
+
     return tensor._from_data(
         tensor._map,
-        aten.squeeze.dims(tensor._value, list(range(1, tensor.dim()))),
+        value,
         tensor._min_key,
         tensor._max_key,
         num_keys=tensor.size(0),
@@ -387,11 +503,14 @@ def _squeeze_dim(
     if tensor._value is None:
         return tensor._shallow_copy()
 
-    dim = [d for d in dim if d != 0 and d != -tensor.dim()]
+    value = tensor.as_tensor()
+    for d in dim[::-1]:
+        if d != 0 and d != -tensor.dim():
+            value = value.squeeze(d)
 
     return tensor._from_data(
         tensor._map,
-        aten.squeeze.dims(tensor._value, dim),
+        value,
         tensor._min_key,
         tensor._max_key,
         num_keys=tensor.size(0),
@@ -406,10 +525,18 @@ def _slice(
     start: Optional[int] = None,
     end: Optional[int] = None,
     step: int = 1,
-) -> Union[HashTensor, Tensor]:
+) -> HashTensor:
 
     if dim == 0 or dim == -tensor.dim():
-        return aten.slice.Tensor(tensor.as_tensor(), dim, start, end, step)
+        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:
+            return tensor._shallow_copy()
+
+        key = aten.slice.Tensor(tensor._key, 0, start, end, step)
+        value = aten.slice.Tensor(tensor.as_tensor(), 0, start, end, step)
+        return tensor.__class__(key, value)
 
     return tensor._from_data(
         tensor._map,
@@ -429,7 +556,7 @@ _old_index_select = torch.index_select
 def _new_index_select(
     input: Tensor,
     dim: int,
-    index: Any,
+    index: Tensor,
     *,
     out: Optional[Tensor] = None,
 ) -> Tensor:
@@ -467,3 +594,54 @@ def _index_select(
         num_keys=tensor.size(0),
         dtype=tensor.dtype,
     )
+
+
+# Since PyTorch does only allow PyTorch tensors as indices in `select`, we need
+# to create a wrapper function and monkey patch `select` :(
+_old_select = torch.select
+
+
+def _new_select(
+    input: Tensor,
+    dim: int,
+    index: int,
+) -> Tensor:
+
+    if 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()):
+        index = int(as_key_tensor([index]))
+    return _old_select(input, dim, index)
+
+
+torch.select = _new_select  # type: ignore
+
+
+@implements(aten.select.int)
+def _select(
+    tensor: HashTensor,
+    dim: int,
+    index: int,
+) -> Union[HashTensor, Tensor]:
+
+    if dim == 0 or dim == -tensor.dim():
+        key = torch.tensor(
+            [index],
+            dtype=tensor._min_key.dtype,
+            device=tensor.device,
+        )
+        return tensor._get(key).squeeze(0)
+
+    return tensor._from_data(
+        tensor._map,
+        aten.select.int(tensor.as_tensor(), dim, index),
+        tensor._min_key,
+        tensor._max_key,
+        num_keys=tensor.size(0),
+        dtype=tensor.dtype,
+    )

{pyg_nightly-2.7.0.dev20250222 → pyg_nightly-2.7.0.dev20250224}/torch_geometric/testing/__init__.py

@@ -22,6 +22,7 @@ from .decorators import (
     withDevice,
     withCUDA,
     withMETIS,
+    withHashTensor,
     disableExtensions,
     withoutExtensions,
 )
@@ -53,6 +54,7 @@ __all__ = [
     'withDevice',
     'withCUDA',
     'withMETIS',
+    'withHashTensor',
     'disableExtensions',
     'withoutExtensions',
     'assert_module',

{pyg_nightly-2.7.0.dev20250222 → pyg_nightly-2.7.0.dev20250224}/torch_geometric/testing/decorators.py

@@ -10,6 +10,7 @@ from packaging.requirements import Requirement
 from packaging.version import Version
 
 import torch_geometric
+import torch_geometric.typing
 from torch_geometric.typing import WITH_METIS, WITH_PYG_LIB, WITH_TORCH_SPARSE
 from torch_geometric.visualization.graph import has_graphviz
 
@@ -265,6 +266,17 @@ def withMETIS(func: Callable) -> Callable:
     )(func)
 
 
+def withHashTensor(func: Callable) -> Callable:
+    r"""A decorator to only test in case :class:`HashTensor` is available."""
+    import pytest
+
+    return pytest.mark.skipif(
+        not torch_geometric.typing.WITH_CPU_HASH_MAP
+        and not has_package('pandas'),
+        reason="HashTensor dependencies not available",
+    )(func)
+
+
 def disableExtensions(func: Callable) -> Callable:
     r"""A decorator to temporarily disable the usage of the
     :obj:`torch_scatter`, :obj:`torch_sparse` and :obj:`pyg_lib` extension