Repository: pyg-team/pytorch_geometric
Pull request #10065 (2025-02-23 20:29:00)

Add more functions to `HashTensor`


Patch:
diff --git a/torch_geometric/hash_tensor.py b/torch_geometric/hash_tensor.py
index 2f02aa897..db894e99e 100644
--- a/torch_geometric/hash_tensor.py
+++ b/torch_geometric/hash_tensor.py
@@ -300,6 +300,7 @@ class HashTensor(Tensor):
         return func(*args, **(kwargs or {}))
 
     def tolist(self) -> List[Any]:
+        """"""  # noqa: D419
         return self.as_tensor().tolist()
 
     def index_select(  # type: ignore
@@ -307,6 +308,7 @@ class HashTensor(Tensor):
         dim: int,
         index: Any,
     ) -> Union['HashTensor', Tensor]:
+        """"""  # noqa: D419
         return torch.index_select(self, dim, index)
 
     def select(  # type: ignore
@@ -314,14 +316,72 @@ 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
+
 
 @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,
@@ -370,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):


Test patch:
diff --git a/test/test_hash_tensor.py b/test/test_hash_tensor.py
index 57c484726..fbaaf7cf9 100644
--- a/test/test_hash_tensor.py
+++ b/test/test_hash_tensor.py
@@ -1,9 +1,13 @@
+import os.path as osp
+from typing import List
+
 import pytest
 import torch
 
 import torch_geometric.typing
 from torch_geometric import HashTensor
-from torch_geometric.testing import withCUDA, withHashTensor
+from torch_geometric.io import fs
+from torch_geometric.testing import onlyCUDA, withCUDA, withHashTensor
 
 KEY_DTYPES = [
     pytest.param(torch.bool, id='bool'),
@@ -72,9 +76,106 @@ def test_empty(dtype, device):
 @withCUDA
 @withHashTensor
 def test_string_key(device):
+    # TODO
     HashTensor(['1', '2', '3'], device=device)
 
 
+@withCUDA
+@withHashTensor
+def test_clone(device):
+    key = torch.tensor([2, 1, 0], device=device)
+    value = torch.randn(key.size(0), 2, device=device)
+    tensor = HashTensor(key, value)
+
+    out = tensor.clone()
+    assert isinstance(out, HashTensor)
+    assert out.dtype == tensor.dtype
+    assert out.device == tensor.device
+    assert out._value.data_ptr() != tensor._value.data_ptr()
+
+    out = torch.clone(tensor)
+    assert isinstance(out, HashTensor)
+    assert out.dtype == tensor.dtype
+    assert out.device == tensor.device
+    assert out._value.data_ptr() != tensor._value.data_ptr()
+
+
+@withCUDA
+@withHashTensor
+def test_share_memory(device):
+    key = torch.tensor([2, 1, 0], device=device)
+    value = torch.randn(key.size(0), 2, device=device)
+    tensor = HashTensor(key, value)
+
+    out = tensor.share_memory_()
+    assert isinstance(out, HashTensor)
+    assert out.is_shared()
+    assert out._value.is_shared()
+    assert out.data_ptr() == tensor.data_ptr()
+
+
+@onlyCUDA
+@withHashTensor
+def test_pin_memory():
+    key = torch.tensor([2, 1, 0])
+    value = torch.randn(key.size(0), 2)
+    tensor = HashTensor(key, value)
+
+    assert not tensor.is_pinned()
+    out = tensor.pin_memory()
+    assert isinstance(out, HashTensor)
+    assert out.is_pinned()
+
+
+@withCUDA
+@withHashTensor
+def test_detach(device):
+    key = torch.tensor([2, 1, 0], device=device)
+    value = torch.randn(key.size(0), 2, device=device, requires_grad=True)
+    tensor = HashTensor(key, value)
+
+    assert tensor.requires_grad
+    out = tensor.detach()
+    assert isinstance(out, HashTensor)
+    assert not out.requires_grad
+    assert not out._value.requires_grad
+
+    tensor.detach_()
+    assert not tensor.requires_grad
+    assert not tensor._value.requires_grad
+
+
+@withCUDA
+@withHashTensor
+def test_contiguous(device):
+    key = torch.tensor([2, 1, 0], device=device)
+    value = torch.randn(2, key.size(0), device=device).t()
+    assert not value.is_contiguous()
+
+    tensor = HashTensor(key, value)
+    assert not tensor.is_contiguous()
+    out = tensor.contiguous()
+    assert out.is_contiguous()
+    assert out._value.is_contiguous()
+
+
+@withCUDA
+@withHashTensor
+def test_save_and_load(device, tmp_path):
+    key = torch.tensor([2, 1, 0], device=device)
+    value = torch.randn(key.size(0), 2, device=device)
+    tensor = HashTensor(key, value)
+
+    path = osp.join(tmp_path, 'hash_tensor.pt')
+    torch.save(tensor, path)
+    out = fs.torch_load(path)
+
+    assert isinstance(out, HashTensor)
+    assert out._value.equal(value)
+    assert out._min_key.equal(key.min())
+    assert out._max_key.equal(key.max())
+
+
 @withCUDA
 @withHashTensor
 def test_to_function(device):
@@ -337,3 +438,33 @@ def test_select(device):
     out = tensor.select(-1, 0)
     assert isinstance(out, HashTensor)
     assert out.as_tensor().equal(value[:, 0])
+
+
+def _collate_fn(hash_tensors: List[HashTensor]) -> List[HashTensor]:
+    return hash_tensors
+
+
+@pytest.mark.parametrize('num_workers', [0, 2])
+@pytest.mark.parametrize('pin_memory', [False, True])
+def test_data_loader(num_workers, pin_memory):
+    key = torch.tensor([2, 1, 0])
+    value = torch.randn(key.size(0), 2)
+    tensor = HashTensor(key, value)
+
+    loader = torch.utils.data.DataLoader(
+        [tensor] * 4,
+        batch_size=2,
+        num_workers=num_workers,
+        collate_fn=_collate_fn,
+        pin_memory=pin_memory,
+        drop_last=True,
+    )
+
+    assert len(loader) == 2
+    for batch in loader:
+        assert isinstance(batch, list)
+        assert len(batch) == 2
+        for tensor in batch:
+            assert isinstance(tensor, HashTensor)
+            assert tensor.dtype == value.dtype
+            assert tensor.is_shared() != (num_workers == 0) or pin_memory