Repository: pyg-team/pytorch_geometric
Pull request #9288 (2024-05-04 17:19:26)

Add `Index[...]` support

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9288?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: Patch coverage is `97.36842%` with `1 lines` in your changes are missing coverage. Please review.
> :exclamation: No coverage uploaded for pull request base (`master@4839cb8`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#section-missing-base-commit).

> :exclamation: Current head b2c20b9 differs from pull request most recent head 2159d12. Consider uploading reports for the commit 2159d12 to get more accurate results

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9288?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Patch % | Lines |
|---|---|---|
| [torch\_geometric/index.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9288?src=pr&el=tree&filepath=torch_geometric%2Findex.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2luZGV4LnB5) | 97.29% | [1 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9288?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff            @@
##             master    #9288   +/-   ##
=========================================
  Coverage          ?   86.97%           
=========================================
  Files             ?      471           
  Lines             ?    30554           
  Branches          ?        0           
=========================================
  Hits              ?    26573           
  Misses            ?     3981           
  Partials          ?        0           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9288?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index bf5e867a6..01f70b23d 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), [#9280](https://github.com/pyg-team/pytorch_geometric/pull/9280), [#9281](https://github.com/pyg-team/pytorch_geometric/pull/9281), [#9284](https://github.com/pyg-team/pytorch_geometric/pull/9284), [#9285](https://github.com/pyg-team/pytorch_geometric/pull/9285), [#9286](https://github.com/pyg-team/pytorch_geometric/pull/9286), [#9287](https://github.com/pyg-team/pytorch_geometric/pull/9287))
+- 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), [#9281](https://github.com/pyg-team/pytorch_geometric/pull/9281), [#9284](https://github.com/pyg-team/pytorch_geometric/pull/9284), [#9285](https://github.com/pyg-team/pytorch_geometric/pull/9285), [#9286](https://github.com/pyg-team/pytorch_geometric/pull/9286), [#9287](https://github.com/pyg-team/pytorch_geometric/pull/9287), [#9288](https://github.com/pyg-team/pytorch_geometric/pull/9288))
 - 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/edge_index.py b/torch_geometric/edge_index.py
index 6c8afc9cc..8b2363dd1 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -1474,14 +1474,14 @@ def _index(
     index = indices[1]
     assert isinstance(index, Tensor)
 
+    out = EdgeIndex(out)
+
     # 1. `edge_index[:, mask]` or `edge_index[..., mask]`.
     if index.dtype in (torch.bool, torch.uint8):
-        out = EdgeIndex(out)
         out._sparse_size = input.sparse_size()
         out._sort_order = input._sort_order
 
     else:  # 2. `edge_index[:, index]` or `edge_index[..., index]`.
-        out = EdgeIndex(out)
         out._sparse_size = input.sparse_size()
 
     return out
diff --git a/torch_geometric/index.py b/torch_geometric/index.py
index 5dfa62ca0..8c11912df 100644
--- a/torch_geometric/index.py
+++ b/torch_geometric/index.py
@@ -424,6 +424,11 @@ def _to_copy(
     )
 
 
+@implements(aten.alias.default)
+def _alias(tensor: Index) -> Index:
+    return tensor._shallow_copy()
+
+
 @implements(aten.sort.default)
 def _sort(
     tensor: Index,
@@ -565,3 +570,31 @@ def _slice(
         out._is_sorted = input.is_sorted
 
     return out
+
+
+@implements(aten.index.Tensor)
+def _index(
+    input: Union[Index, Tensor],
+    indices: List[Optional[Union[Tensor, Index]]],
+) -> Union[Index, Tensor]:
+
+    if not isinstance(input, Index):
+        indices = pytree.tree_map_only(Index, lambda x: x._data, indices)
+        return aten.index.Tensor(input, indices)
+
+    data = aten.index.Tensor(input._data, indices)
+
+    assert len(indices) == 1
+    index = indices[0]
+    assert index is not None
+
+    out = Index(data)
+
+    if index.dtype in (torch.bool, torch.uint8):  # 1. `index[mask]`.
+        out._dim_size = input.dim_size
+        out._is_sorted = input.is_sorted
+
+    else:  # 2. `index[index]`.
+        out._dim_size = input.dim_size
+
+    return out


Test patch:
diff --git a/test/test_index.py b/test/test_index.py
index 2c07b73bc..3692105af 100644
--- a/test/test_index.py
+++ b/test/test_index.py
@@ -352,3 +352,65 @@ def test_narrow(dtype, device):
     assert out.equal(tensor([1, 1], device=device))
     assert out.dim_size == 3
     assert out.is_sorted
+
+
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_getitem(dtype, device):
+    kwargs = dict(dtype=dtype, device=device)
+    index = Index([0, 1, 1, 2], dim_size=3, is_sorted=True, **kwargs)
+
+    out = index[:]
+    assert isinstance(out, Index)
+    assert out._data.data_ptr() == index._data.data_ptr()
+    assert out.equal(tensor([0, 1, 1, 2], device=device))
+    assert out.dim_size == 3
+    assert out.is_sorted
+
+    out = index[torch.tensor([False, True, False, True], device=device)]
+    assert isinstance(out, Index)
+    assert out.equal(tensor([1, 2], device=device))
+    assert out.dim_size == 3
+    assert out.is_sorted
+
+    out = index[tensor([1, 3], device=device)]
+    assert isinstance(out, Index)
+    assert out.equal(tensor([1, 2], device=device))
+    assert out.dim_size == 3
+    assert not out.is_sorted
+
+    out = index[1:3]
+    assert isinstance(out, Index)
+    assert out.equal(tensor([1, 1], device=device))
+    assert out.dim_size == 3
+    assert out.is_sorted
+
+    out = index[...]
+    assert isinstance(out, Index)
+    assert out._data.data_ptr() == index._data.data_ptr()
+    assert out.equal(tensor([0, 1, 1, 2], device=device))
+    assert out.dim_size == 3
+    assert out.is_sorted
+
+    out = index[..., 1:3]
+    assert isinstance(out, Index)
+    assert out.equal(tensor([1, 1], device=device))
+    assert out.dim_size == 3
+    assert out.is_sorted
+
+    out = index[None, 1:3]
+    assert not isinstance(out, Index)
+    assert out.equal(tensor([[1, 1]], device=device))
+
+    out = index[1:3, None]
+    assert not isinstance(out, Index)
+    assert out.equal(tensor([[1], [1]], device=device))
+
+    out = index[0]
+    assert not isinstance(out, Index)
+    assert out.equal(tensor(0, device=device))
+
+    tmp = torch.randn(3, device=device)
+    out = tmp[index]
+    assert not isinstance(out, Index)
+    assert out.equal(tmp[index.as_tensor()])