Repository: pyg-team/pytorch_geometric
Pull request #9277 (2024-05-04 10:34:33)

Add `torch_geometric.Index` tests


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 75826e1cc..38edaa119 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))
+- Added `torch_geometric.Index` ([#9276](https://github.com/pyg-team/pytorch_geometric/pull/9276), [#9277](https://github.com/pyg-team/pytorch_geometric/pull/9277))
 - 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/__init__.py b/torch_geometric/__init__.py
index 29a636833..381c3524e 100644
--- a/torch_geometric/__init__.py
+++ b/torch_geometric/__init__.py
@@ -1,4 +1,5 @@
 from ._compile import compile, is_compiling
+from .index import Index
 from .edge_index import EdgeIndex
 from .seed import seed_everything
 from .home import get_home_dir, set_home_dir
@@ -25,6 +26,7 @@ graphgym = LazyLoader('graphgym', globals(), 'torch_geometric.graphgym')
 __version__ = '2.6.0'
 
 __all__ = [
+    'Index',
     'EdgeIndex',
     'seed_everything',
     'get_home_dir',
diff --git a/torch_geometric/index.py b/torch_geometric/index.py
index a0b013186..2382ec2ee 100644
--- a/torch_geometric/index.py
+++ b/torch_geometric/index.py
@@ -1,7 +1,19 @@
 import functools
-from typing import Any, Callable, Dict, List, NamedTuple, Optional, Type, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    NamedTuple,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+)
 
 import torch
+import torch.utils._pytree as pytree
 from torch import Tensor
 
 from torch_geometric.typing import INDEX_DTYPES
@@ -118,7 +130,7 @@ class Index(Tensor):
 
         if isinstance(data, cls):  # If passed `Index`, inherit metadata:
             indptr = data._indptr
-            dim_size = dim_size or data.dim_size()
+            dim_size = dim_size or data.dim_size
             is_sorted = is_sorted or data.is_sorted
 
         assert_valid_dtype(data)
@@ -146,7 +158,7 @@ class Index(Tensor):
             out._data = data._data
 
             # Reset metadata if cache is invalidated:
-            if dim_size is not None and dim_size != data.dim_size():
+            if dim_size is not None and dim_size != data.dim_size:
                 out._indptr = None
 
         return out
@@ -154,31 +166,38 @@ class Index(Tensor):
     # Validation ##############################################################
 
     def validate(self) -> 'Index':
-        raise NotImplementedError
+        r"""TODO."""
+        return self
 
     # Properties ##############################################################
 
     @property
     def dim_size(self) -> Optional[int]:
-        raise NotImplementedError
+        r"""TODO."""
+        return self._dim_size
 
     @property
     def is_sorted(self) -> bool:
-        raise NotImplementedError
+        r"""TODO."""
+        return self._is_sorted
 
     # Cache Interface #########################################################
 
     def get_dim_size(self) -> int:
+        r"""TODO."""
         raise NotImplementedError
 
     def dim_resize_(self, dim_size: Optional[int]) -> 'Index':
+        r"""TODO."""
         raise NotImplementedError
 
     @assert_sorted
     def get_indptr(self) -> Tensor:
+        r"""TODO."""
         raise NotImplementedError
 
     def fill_cache_(self) -> 'Index':
+        r"""TODO."""
         raise NotImplementedError
 
     # Methods #################################################################
@@ -199,8 +218,93 @@ class Index(Tensor):
         return self._data
 
     def dim_narrow(self, start: Union[int, Tensor], length: int) -> 'Index':
+        r"""TODO."""
         raise NotImplementedError
 
+    # PyTorch/Python builtins #################################################
+
+    def __tensor_flatten__(self) -> Tuple[List[str], Tuple[Any, ...]]:
+        attrs = ['_data']
+        if self._indptr is not None:
+            attrs.append('_indptr')
+
+        ctx = (
+            self._dim_size,
+            self._is_sorted,
+            self._cat_metadata,
+        )
+
+        return attrs, ctx
+
+    @staticmethod
+    def __tensor_unflatten__(
+        inner_tensors: Dict[str, Any],
+        ctx: Tuple[Any, ...],
+        outer_size: Tuple[int, ...],
+        outer_stride: Tuple[int, ...],
+    ) -> 'Index':
+        index = Index(
+            inner_tensors['_data'],
+            dim_size=ctx[0],
+            is_sorted=ctx[1],
+        )
+
+        index._indptr = inner_tensors.get('_indptr', None)
+        index._cat_metadata = ctx[2]
+
+        return index
+
+    # Prevent auto-wrapping outputs back into the proper subclass type:
+    __torch_function__ = torch._C._disabled_torch_function_impl
+
+    @classmethod
+    def __torch_dispatch__(
+        cls: Type,
+        func: Callable[..., Any],
+        types: Iterable[Type[Any]],
+        args: Iterable[Tuple[Any, ...]] = (),
+        kwargs: Optional[Dict[Any, Any]] = None,
+    ) -> Any:
+        # `Index` should be treated as a regular PyTorch tensor for all
+        # standard PyTorch functionalities. However,
+        # * some of its metadata can be transferred to new functions, e.g.,
+        #   `torch.narrow()` can inherit the `is_sorted` property.
+        # * not all operations lead to valid `Index` tensors again, e.g.,
+        #   `torch.sum()` does not yield a `Index` as its output, or
+        #   `torch.stack() violates the [*] shape assumption.
+
+        # To account for this, we hold a number of `HANDLED_FUNCTIONS` that
+        # implement specific functions for valid `Index` 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(Index, lambda x: x._data, args)
+        if kwargs is not None:
+            kwargs = pytree.tree_map_only(Index, lambda x: x._data, kwargs)
+        return func(*args, **(kwargs or {}))
+
+    def __repr__(self) -> str:  # type: ignore
+        prefix = f'{self.__class__.__name__}('
+        indent = len(prefix)
+        tensor_str = torch._tensor_str._tensor_str(self._data, indent)
+
+        suffixes = []
+        if self.dim_size is not None:
+            suffixes.append(f'dim_size={self.dim_size}')
+        if (self.device.type != torch._C._get_default_device()
+                or (self.device.type == 'cuda'
+                    and torch.cuda.current_device() != self.device.index)
+                or (self.device.type == 'mps')):
+            suffixes.append(f"device='{self.device}'")
+        if self.dtype != torch.int64:
+            suffixes.append(f'dtype={self.dtype}')
+        if self.is_sorted:
+            suffixes.append('is_sorted=True')
+
+        return torch._tensor_str._add_suffixes(prefix + tensor_str, suffixes,
+                                               indent, force_newline=False)
+
 
 # def sort(self) -> None:
 #     # TODO MOVE BEHIND TORCH DISPATCH


Test patch:
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index d9fdb404a..e7f82fe92 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -46,12 +46,8 @@ def test_basic(dtype, device):
     adj.validate()
     assert isinstance(adj, EdgeIndex)
 
-    if torch_geometric.typing.WITH_PT112:
-        assert str(adj).startswith('EdgeIndex([[0, 1, 1, 2],\n'
-                                   '           [1, 0, 2, 1]], ')
-    else:
-        assert str(adj).startswith('tensor([[0, 1, 1, 2],\n'
-                                   '        [1, 0, 2, 1]], ')
+    assert str(adj).startswith('EdgeIndex([[0, 1, 1, 2],\n'
+                               '           [1, 0, 2, 1]], ')
     assert 'sparse_size=(3, 3), nnz=4' in str(adj)
     assert (f"device='{device}'" in str(adj)) == adj.is_cuda
     assert (f'dtype={dtype}' in str(adj)) == (dtype != torch.long)
diff --git a/test/test_index.py b/test/test_index.py
new file mode 100644
index 000000000..c9928a4a7
--- /dev/null
+++ b/test/test_index.py
@@ -0,0 +1,56 @@
+import pytest
+import torch
+
+from torch_geometric import Index
+from torch_geometric.testing import withCUDA
+from torch_geometric.typing import INDEX_DTYPES
+
+DTYPES = [pytest.param(dtype, id=str(dtype)[6:]) for dtype in INDEX_DTYPES]
+
+
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_basic(dtype, device):
+    kwargs = dict(dtype=dtype, device=device, dim_size=3)
+    index = Index([0, 1, 1, 2], **kwargs)
+    index.validate()
+    assert isinstance(index, Index)
+
+    assert str(index).startswith('Index([0, 1, 1, 2], ')
+    assert 'dim_size=3' in str(index)
+    assert (f"device='{device}'" in str(index)) == index.is_cuda
+    assert (f'dtype={dtype}' in str(index)) == (dtype != torch.long)
+
+    assert index.dtype == dtype
+    assert index.device == device
+    assert index.dim_size == 3
+    assert not index.is_sorted
+
+    out = index.as_tensor()
+    assert not isinstance(out, Index)
+    assert out.dtype == dtype
+    assert out.device == device
+
+    out = index * 1
+    assert not isinstance(out, Index)
+    assert out.dtype == dtype
+    assert out.device == device
+
+
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_identity(dtype, device):
+    kwargs = dict(dtype=dtype, device=device, dim_size=3, is_sorted=True)
+    index = Index([0, 1, 1, 2], **kwargs)
+
+    out = Index(index)
+    assert not isinstance(out.as_tensor(), Index)
+    assert out.data_ptr() == index.data_ptr()
+    assert out.dtype == index.dtype
+    assert out.device == index.device
+    assert out.dim_size == index.dim_size
+    assert out.is_sorted == index.is_sorted
+
+    out = Index(index, dim_size=4, is_sorted=False)
+    assert out.dim_size == 4
+    assert out.is_sorted == index.is_sorted