Repository: pyg-team/pytorch_geometric
Pull request #8455 (2023-11-27 21:41:37)

`EdgeIndex.__getitem__` support


Patch:
diff --git a/torch_geometric/data/edge_index.py b/torch_geometric/data/edge_index.py
index 3f6ed2958..36336cc31 100644
--- a/torch_geometric/data/edge_index.py
+++ b/torch_geometric/data/edge_index.py
@@ -1,6 +1,6 @@
 import functools
 from enum import Enum
-from typing import Callable, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -420,3 +420,44 @@ def narrow(
         out._sort_order = input._sort_order
 
     return out
+
+
+@implements(Tensor.__getitem__)
+def getitem(input: EdgeIndex, index: Any) -> Union[EdgeIndex, Tensor]:
+    out = Tensor.__torch_function__(  #
+        Tensor.__getitem__, (Tensor, ), (input, index))
+
+    # There exists 3 possible index types that map back to a valid `EdgeIndex`,
+    # and all include selecting/filtering in the last dimension only:
+    def is_last_dim_select(i: Any) -> bool:
+        # Maps to true for `__getitem__` requests of the form
+        # `tensor[..., index]` or `tensor[:, index]`.
+        if not isinstance(i, tuple) or len(i) != 2:
+            return False
+        if i[0] == Ellipsis:
+            return True
+        if not isinstance(i[0], slice):
+            return False
+        return i[0].start is None and i[0].stop is None and i[0].step is None
+
+    is_valid = is_last_dim_select(index)
+
+    # 1. `edge_index[:, mask]` or `edge_index[..., mask]`.
+    if is_valid and isinstance(index[1], (torch.BoolTensor, torch.ByteTensor)):
+        out = out.as_subclass(EdgeIndex)
+        out._sparse_size = input.sparse_size
+        out._sort_order = input._sort_order
+
+    # 2. `edge_index[:, index]` or `edge_index[..., index]`.
+    elif is_valid and isinstance(index[1], Tensor):
+        out = out.as_subclass(EdgeIndex)
+        out._sparse_size = input.sparse_size
+
+    # 3. `edge_index[:, slice]` or `edge_index[..., slice]`.
+    elif is_valid and isinstance(index[1], slice):
+        out = out.as_subclass(EdgeIndex)
+        out._sparse_size = input.sparse_size
+        if index[1].step is None or index[1].step > 0:
+            out._sort_order = input._sort_order
+
+    return out


Test patch:
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
index 0ccc81d41..518269851 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -62,6 +62,16 @@ def test_cpu_cuda():
     assert not out.is_cuda
 
 
+def test_share_memory():
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
+    adj.fill_cache()
+
+    adj = adj.share_memory_()
+    assert isinstance(adj, EdgeIndex)
+    assert adj.is_shared()
+    assert adj._rowptr.is_shared()
+
+
 def test_sort_by():
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
     out = adj.sort_by('row')
@@ -166,6 +176,31 @@ def test_narrow():
     assert not isinstance(out, EdgeIndex)
 
 
+def test_getitem():
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
+
+    out = adj[:, torch.tensor([False, True, False, True])]
+    assert isinstance(out, EdgeIndex)
+    assert torch.equal(out, torch.tensor([[1, 2], [0, 1]]))
+    assert out.sort_order == 'row'
+
+    out = adj[..., torch.tensor([1, 3])]
+    assert isinstance(out, EdgeIndex)
+    assert torch.equal(out, torch.tensor([[1, 2], [0, 1]]))
+    assert out.sort_order is None
+
+    out = adj[..., 1::2]
+    assert isinstance(out, EdgeIndex)
+    assert torch.equal(out, torch.tensor([[1, 2], [0, 1]]))
+    assert out.sort_order == 'row'
+
+    out = adj[:, 0]
+    assert not isinstance(out, EdgeIndex)
+
+    out = adj[torch.tensor([0])]
+    assert not isinstance(out, EdgeIndex)
+
+
 def test_save_and_load(tmp_path):
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
     adj.fill_cache()
@@ -183,16 +218,6 @@ def test_save_and_load(tmp_path):
     assert torch.equal(out._rowptr, torch.tensor([0, 1, 3, 4]))
 
 
-def test_share_memory():
-    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
-    adj.fill_cache()
-
-    adj = adj.share_memory_()
-    assert isinstance(adj, EdgeIndex)
-    assert adj.is_shared()
-    assert adj._rowptr.is_shared()
-
-
 @pytest.mark.parametrize('num_workers', [0, 2])
 def test_data_loader(num_workers):
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')