Repository: pyg-team/pytorch_geometric
Pull request #9286 (2024-05-04 16:46:10)

Add `Index.index_select()`


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 609fd587c..1d0265f3a 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))
+- 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))
 - 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 ac762bc5d..2298fb79f 100644
--- a/torch_geometric/index.py
+++ b/torch_geometric/index.py
@@ -511,3 +511,18 @@ def _flip(
     out._dim_size = input.dim_size
 
     return out
+
+
+@implements(aten.index_select.default)
+def _index_select(
+    input: Index,
+    dim: int,
+    index: Tensor,
+) -> Index:
+
+    data = aten.index_select.default(input._data, dim, index)
+
+    out = Index(data)
+    out._dim_size = input.dim_size
+
+    return out


Test patch:
diff --git a/test/test_index.py b/test/test_index.py
index 57bc1c9bc..48e8bca02 100644
--- a/test/test_index.py
+++ b/test/test_index.py
@@ -318,3 +318,24 @@ def test_flip(dtype, device):
     assert out.equal(tensor([2, 1, 1, 0], device=device))
     assert out.dim_size == 3
     assert not out.is_sorted
+
+
+@withCUDA
+@pytest.mark.parametrize('dtype', DTYPES)
+def test_index_select(dtype, device):
+    kwargs = dict(dtype=dtype, device=device)
+    index = Index([0, 1, 1, 2], dim_size=3, is_sorted=True, **kwargs)
+
+    i = tensor([1, 3], device=device)
+    out = index.index_select(0, i)
+    assert out.equal(tensor([1, 2], device=device))
+    assert isinstance(out, Index)
+    assert out.dim_size == 3
+    assert not out.is_sorted
+
+    inplace = torch.empty(2, dtype=dtype, device=device)
+    out = torch.index_select(index, 0, i, out=inplace)
+    assert out.equal(tensor([1, 2], device=device))
+    assert out.data_ptr() == inplace.data_ptr()
+    assert not isinstance(out, Index)
+    assert not isinstance(inplace, Index)