Repository: pyg-team/pytorch_geometric
Pull request #10031 (2025-02-14 20:13:44)

Add `HashTensor` boilerplate


Patch:
diff --git a/torch_geometric/__init__.py b/torch_geometric/__init__.py
index c65c07cd3..1bab1f75c 100644
--- a/torch_geometric/__init__.py
+++ b/torch_geometric/__init__.py
@@ -7,6 +7,7 @@ from ._compile import compile, is_compiling
 from ._onnx import is_in_onnx_export
 from .index import Index
 from .edge_index import EdgeIndex
+from .hash_tensor import HashTensor
 from .seed import seed_everything
 from .home import get_home_dir, set_home_dir
 from .device import is_mps_available, is_xpu_available, device
@@ -35,6 +36,7 @@ __version__ = '2.7.0'
 __all__ = [
     'Index',
     'EdgeIndex',
+    'HashTensor',
     'seed_everything',
     'get_home_dir',
     'set_home_dir',
@@ -67,4 +69,5 @@ if torch_geometric.typing.WITH_PT24:
         EdgeIndex,
         torch_geometric.edge_index.SortOrder,
         torch_geometric.edge_index.CatMetadata,
+        HashTensor,
     ])
diff --git a/torch_geometric/hash_tensor.py b/torch_geometric/hash_tensor.py
new file mode 100644
index 000000000..d4407851e
--- /dev/null
+++ b/torch_geometric/hash_tensor.py
@@ -0,0 +1,112 @@
+from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Type, Union
+
+import torch
+import torch.utils._pytree as pytree
+from torch import Tensor
+
+from torch_geometric.typing import CPUHashMap, CUDAHashMap
+
+HANDLED_FUNCTIONS: Dict[Callable, Callable] = {}
+
+
+class HashTensor(Tensor):
+    _map: Union[Tensor, CPUHashMap, CUDAHashMap]
+    _value: Optional[Tensor]
+
+    @staticmethod
+    def __new__(
+        cls: Type,
+        key: Any,
+        value: Optional[Any] = None,
+        *,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+    ) -> 'HashTensor':
+
+        if value is not None:
+            if not isinstance(value, Tensor):
+                value = torch.tensor(value, dtype=dtype, device=device)
+            else:
+                value = value.to(device, dtype)
+
+        # TODO Add numpy support
+        if not isinstance(key, Tensor):
+            device = value.device if value is not None else device
+            key = torch.tensor(key, device=device)
+        else:
+            key = key.to(device)
+            device = key.device
+
+        if key.dim() != 1:
+            raise ValueError(f"'key' data in '{cls.__name__}' needs to be "
+                             f"one-dimensional (got {key.dim()} dimensions)")
+
+        if not key.is_contiguous():
+            raise ValueError(f"'key' data in '{cls.__name__}' needs to be "
+                             f"contiguous")
+
+        if value is not None:
+            if key.device != value.device:
+                raise ValueError(f"'key' and 'value' data in '{cls.__name__}' "
+                                 f"are expected to be on the same device (got "
+                                 f"'{key.device}' and '{value.device}')")
+
+            if key.numel() != value.size(0):
+                raise ValueError(f"'key' and 'value' data in '{cls.__name__}' "
+                                 f"are expected to have the same size in the "
+                                 f"first dimension (got {key.size(0)} and "
+                                 f"{value.size(0)})")
+
+            dtype = value.dtype
+            size = value.size()
+            stride = value.stride()
+            layout = value.layout
+            requires_grad = value.requires_grad
+        else:
+            dtype = dtype or torch.int64
+            size = torch.Size([key.numel()])
+            stride = (1, )
+            layout = torch.strided
+            requires_grad = False
+
+        out = Tensor._make_wrapper_subclass(  # type: ignore
+            cls,
+            size=size,
+            strides=stride,
+            dtype=dtype,
+            device=device,
+            layout=layout,
+            requires_grad=requires_grad,
+        )
+
+        out._value = value
+
+        return out
+
+    def as_tensor(self) -> Tensor:
+        r"""Zero-copies the :class:`HashTensor` representation back to a
+        :class:`torch.Tensor` representation.
+        """
+        if self._value is not None:
+            return self._value
+        return torch.arange(self.size(0), dtype=self.dtype, device=self.device)
+
+    @classmethod
+    def __torch_dispatch__(
+        cls: Type,
+        func: Callable[..., Any],
+        types: Iterable[Type[Any]],
+        args: Iterable[Tuple[Any, ...]] = (),
+        kwargs: Optional[Dict[Any, Any]] = None,
+    ) -> Any:
+        # Hold a number of `HANDLED_FUNCTIONS` that implement specific
+        # functions for valid `HashTensor` routines.
+        if func in HANDLED_FUNCTIONS:
+            return HANDLED_FUNCTIONS[func](*args, **(kwargs or {}))
+
+        # For all other PyTorch functions, we treat them as vanilla tensors.
+        args = pytree.tree_map_only(HashTensor, lambda x: x.as_tensor(), args)
+        if kwargs is not None:
+            kwargs = pytree.tree_map_only(HashTensor, lambda x: x.as_tensor(),
+                                          kwargs)
+        return func(*args, **(kwargs or {}))
diff --git a/torch_geometric/typing.py b/torch_geometric/typing.py
index 513f04184..e24d2cb12 100644
--- a/torch_geometric/typing.py
+++ b/torch_geometric/typing.py
@@ -9,6 +9,11 @@ import numpy as np
 import torch
 from torch import Tensor
 
+try:
+    from typing import TypeAlias  # type: ignore
+except ImportError:
+    from typing_extensions import TypeAlias
+
 WITH_PT20 = int(torch.__version__.split('.')[0]) >= 2
 WITH_PT21 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 1
 WITH_PT22 = WITH_PT20 and int(torch.__version__.split('.')[1]) >= 2
@@ -64,6 +69,10 @@ try:
         pyg_lib.sampler.neighbor_sample).parameters)
     WITH_WEIGHTED_NEIGHBOR_SAMPLE = ('edge_weight' in inspect.signature(
         pyg_lib.sampler.neighbor_sample).parameters)
+    # WITH_CPU_HASH_MAP = hasattr(torch.classes.pyg, 'CPUHashMap')
+    # WITH_CUDA_HASH_MAP = hasattr(torch.classes.pyg, 'CUDAHashMap')
+    WITH_CPU_HASH_MAP = False
+    WITH_CUDA_HASH_MAP = False
 except Exception as e:
     if not isinstance(e, ImportError):  # pragma: no cover
         warnings.warn(f"An issue occurred while importing 'pyg-lib'. "
@@ -78,6 +87,32 @@ except Exception as e:
     WITH_METIS = False
     WITH_EDGE_TIME_NEIGHBOR_SAMPLE = False
     WITH_WEIGHTED_NEIGHBOR_SAMPLE = False
+    WITH_CPU_HASH_MAP = False
+    WITH_CUDA_HASH_MAP = False
+
+if WITH_CPU_HASH_MAP:
+    CPUHashMap: TypeAlias = torch.classes.pyg.CPUHashMap
+else:
+
+    class CPUHashMap:  # type: ignore
+        def __init__(self, key: Tensor) -> None:
+            raise ImportError("'CPUHashMap' requires 'pyg-lib'")
+
+        def get(self, query: Tensor) -> Tensor:
+            raise ImportError("'CPUHashMap' requires 'pyg-lib'")
+
+
+if WITH_CUDA_HASH_MAP:
+    CUDAHashMap: TypeAlias = torch.classes.pyg.CUDAHashMap
+else:
+
+    class CUDAHashMap:  # type: ignore
+        def __init__(self, key: Tensor) -> None:
+            raise ImportError("'CUDAHashMap' requires 'pyg-lib'")
+
+        def get(self, query: Tensor) -> Tensor:
+            raise ImportError("'CUDAHashMap' requires 'pyg-lib'")
+
 
 try:
     import torch_scatter  # noqa


Test patch:
diff --git a/test/test_hash_tensor.py b/test/test_hash_tensor.py
new file mode 100644
index 000000000..3e8ba9bcc
--- /dev/null
+++ b/test/test_hash_tensor.py
@@ -0,0 +1,30 @@
+import pytest
+import torch
+
+from torch_geometric import HashTensor
+from torch_geometric.testing import withCUDA
+
+KEY_DTYPES = [
+    torch.bool,
+    torch.uint8,
+    torch.int8,
+    torch.int16,
+    torch.int32,
+    torch.int64,
+    torch.float16,
+    torch.bfloat16,
+    torch.float32,
+    torch.float64,
+]
+
+
+@withCUDA
+@pytest.mark.parametrize('dtype', KEY_DTYPES)
+def test_basic(dtype, device):
+    if dtype != torch.bool:
+        key = torch.tensor([2, 1, 0], dtype=dtype, device=device)
+    else:
+        key = torch.tensor([True, False], device=device)
+    value = torch.randn(key.size(0), 2, device=device)
+
+    HashTensor(key, value)