Repository: pyg-team/pytorch_geometric
Pull request #9280 (2024-05-04 12:46:04)

Add common tensor utilities to `torch_geometric.Index`


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2d2e853b2..5dedb1449 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,7 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
-- Added `torch_geometric.Index` ([#9276](https://github.com/pyg-team/pytorch_geometric/pull/9276), [#9277](https://github.com/pyg-team/pytorch_geometric/pull/9277), [#9278](https://github.com/pyg-team/pytorch_geometric/pull/9278), [#9279](https://github.com/pyg-team/pytorch_geometric/pull/9279))
+- Added `torch_geometric.Index` ([#9276](https://github.com/pyg-team/pytorch_geometric/pull/9276), [#9277](https://github.com/pyg-team/pytorch_geometric/pull/9277), [#9278](https://github.com/pyg-team/pytorch_geometric/pull/9278), [#9279](https://github.com/pyg-team/pytorch_geometric/pull/9279), [#9280](https://github.com/pyg-team/pytorch_geometric/pull/9280))
 - Added support for PyTorch 2.3 ([#9240](https://github.com/pyg-team/pytorch_geometric/pull/9240))
 - Added support for `EdgeIndex` in `message_and_aggregate` ([#9131](https://github.com/pyg-team/pytorch_geometric/pull/9131))
 - Added `CornellTemporalHyperGraphDataset` ([#9090](https://github.com/pyg-team/pytorch_geometric/pull/9090))
diff --git a/torch_geometric/index.py b/torch_geometric/index.py
index 32f5b2931..c27e19e98 100644
--- a/torch_geometric/index.py
+++ b/torch_geometric/index.py
@@ -9,6 +9,7 @@ from typing import (
     Optional,
     Tuple,
     Type,
+    Union,
 )
 
 import torch
@@ -349,6 +350,70 @@ class Index(Tensor):
                                                indent, force_newline=False)
 
 
+def apply_(
+    tensor: Index,
+    fn: Callable,
+    *args: Any,
+    **kwargs: Any,
+) -> Union[Index, Tensor]:
+
+    data = fn(tensor._data, *args, **kwargs)
+
+    if data.dtype not in INDEX_DTYPES:
+        return data
+
+    if tensor._data.data_ptr() != data.data_ptr():
+        out = Index(data)
+    else:  # In-place:
+        tensor._data = data
+        out = tensor
+
+    # Copy metadata:
+    out._dim_size = tensor._dim_size
+    out._is_sorted = tensor._is_sorted
+    out._cat_metadata = tensor._cat_metadata
+
+    # Convert cache:
+    if tensor._indptr is not None:
+        out._indptr = fn(tensor._indptr, *args, **kwargs)
+
+    return out
+
+
+@implements(aten.clone.default)
+def _clone(
+    tensor: Index,
+    *,
+    memory_format: torch.memory_format = torch.preserve_format,
+) -> Index:
+    out = apply_(tensor, aten.clone.default, memory_format=memory_format)
+    assert isinstance(out, Index)
+    return out
+
+
+@implements(aten._to_copy.default)
+def _to_copy(
+    tensor: Index,
+    *,
+    dtype: Optional[torch.dtype] = None,
+    layout: Optional[torch.layout] = None,
+    device: Optional[torch.device] = None,
+    pin_memory: bool = False,
+    non_blocking: bool = False,
+    memory_format: Optional[torch.memory_format] = None,
+) -> Union[Index, Tensor]:
+    return apply_(
+        tensor,
+        aten._to_copy.default,
+        dtype=dtype,
+        layout=layout,
+        device=device,
+        pin_memory=pin_memory,
+        non_blocking=non_blocking,
+        memory_format=memory_format,
+    )
+
+
 # def sort(self) -> None:
 #     # TODO MOVE BEHIND TORCH DISPATCH
 #     raise NotImplementedError


Test patch:
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index 596006c5f..26b7ed484 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -301,6 +301,7 @@ def test_to_function(dtype, device, is_undirected):
 def test_cpu_cuda(dtype):
     kwargs = dict(dtype=dtype)
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], **kwargs)
+    assert adj.is_cpu
 
     out = adj.cuda()
     assert isinstance(out, EdgeIndex)
@@ -308,7 +309,7 @@ def test_cpu_cuda(dtype):
 
     out = out.cpu()
     assert isinstance(out, EdgeIndex)
-    assert not out.is_cuda
+    assert out.is_cpu
 
 
 @withCUDA
diff --git a/test/test_index.py b/test/test_index.py
index 28cdb09aa..20b828adf 100644
--- a/test/test_index.py
+++ b/test/test_index.py
@@ -2,8 +2,9 @@ import pytest
 import torch
 from torch import tensor
 
+import torch_geometric.typing
 from torch_geometric import Index
-from torch_geometric.testing import withCUDA
+from torch_geometric.testing import onlyCUDA, withCUDA
 from torch_geometric.typing import INDEX_DTYPES
 
 DTYPES = [pytest.param(dtype, id=str(dtype)[6:]) for dtype in INDEX_DTYPES]
@@ -110,3 +111,108 @@ def test_dim_resize(dtype, device):
     out = index.dim_resize_(None)
     assert out.dim_size is None
     assert out._indptr is None
+
+
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_clone(dtype, device):
+    kwargs = dict(dtype=dtype, device=device)
+    index = Index([0, 1, 1, 2], is_sorted=True, dim_size=3, **kwargs)
+
+    out = index.clone()
+    assert isinstance(out, Index)
+    assert out.dtype == dtype
+    assert out.device == device
+    assert out.dim_size == 3
+    assert out.is_sorted
+
+    out = torch.clone(index)
+    assert isinstance(out, Index)
+    assert out.dtype == dtype
+    assert out.device == device
+    assert out.dim_size == 3
+    assert out.is_sorted
+
+
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_to_function(dtype, device):
+    kwargs = dict(dtype=dtype)
+    index = Index([0, 1, 1, 2], dim_size=3, is_sorted=True, **kwargs)
+    index.fill_cache_()
+
+    index = index.to(device)
+    assert isinstance(index, Index)
+    assert index.device == device
+    assert index._indptr.dtype == dtype
+    assert index._indptr.device == device
+
+    out = index.cpu()
+    assert isinstance(out, Index)
+    assert out.device == torch.device('cpu')
+
+    out = index.to(torch.int)
+    assert out.dtype == torch.int
+    if torch_geometric.typing.WITH_PT20:
+        assert isinstance(out, Index)
+        assert out._indptr.dtype == torch.int
+    else:
+        assert not isinstance(out, Index)
+
+    out = index.to(torch.float)
+    assert not isinstance(out, Index)
+    assert out.dtype == torch.float
+
+    out = index.long()
+    assert isinstance(out, Index)
+    assert out.dtype == torch.int64
+
+    out = index.int()
+    assert out.dtype == torch.int
+    if torch_geometric.typing.WITH_PT20:
+        assert isinstance(out, Index)
+    else:
+        assert not isinstance(out, Index)
+
+
+@onlyCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_cpu_cuda(dtype):
+    kwargs = dict(dtype=dtype)
+    index = Index([0, 1, 1, 2], dim_size=3, is_sorted=True, **kwargs)
+    assert index.is_cpu
+
+    out = index.cuda()
+    assert isinstance(out, Index)
+    assert out.is_cuda
+
+    out = out.cpu()
+    assert isinstance(out, Index)
+    assert out.is_cpu
+
+
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_share_memory(dtype, device):
+    kwargs = dict(dtype=dtype, device=device)
+    index = Index([0, 1, 1, 2], dim_size=3, is_sorted=True, **kwargs)
+    index.fill_cache_()
+
+    out = index.share_memory_()
+    assert isinstance(out, Index)
+    assert out.is_shared()
+    assert out._data.is_shared()
+    assert out._indptr.is_shared()
+    assert out.data_ptr() == index.data_ptr()
+
+
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_contiguous(dtype, device):
+    kwargs = dict(dtype=dtype, device=device)
+    index = Index([0, 1, 1, 2], dim_size=3, is_sorted=True, **kwargs)
+
+    assert index.is_contiguous
+    out = index.contiguous()
+    assert isinstance(out, Index)
+    assert out.is_contiguous