Repository: pyg-team/pytorch_geometric
Pull request #8450 (2023-11-27 13:35:54)

`EdgeIndex.index_select` and `EdgeIndex.narrow`


Patch:
diff --git a/torch_geometric/data/edge_index.py b/torch_geometric/data/edge_index.py
index 630bf7d4a..0df14e3bd 100644
--- a/torch_geometric/data/edge_index.py
+++ b/torch_geometric/data/edge_index.py
@@ -228,7 +228,7 @@ class EdgeIndex(Tensor):
         out = self.__class__(edge_index)
 
         # We can fully fill metadata and cache:
-        out._sparse_size = self._sparse_size
+        out._sparse_size = self.sparse_size
         out._sort_order = sort_order
         out._rowptr = self._rowptr
         out._colptr = self._colptr
@@ -275,7 +275,7 @@ def apply_(
     out = out.as_subclass(EdgeIndex)
 
     # Copy metadata:
-    out._sparse_size = out._sparse_size
+    out._sparse_size = out.sparse_size
     out._sort_order = out._sort_order
 
     # Convert cache:
@@ -360,7 +360,7 @@ def flip(
 
     # Flip metadata and cache:
     if 0 in dims or -2 in dims:
-        out._sparse_size = input._sparse_size[::-1]
+        out._sparse_size = input.sparse_size[::-1]
 
     if len(dims) == 1 and (dims[0] == 0 or dims[0] == -2):
         if input._sort_order == SortOrder.ROW:
@@ -374,3 +374,44 @@ def flip(
         out._csc2csr = input._csr2csc
 
     return out
+
+
+@implements(torch.index_select)
+@implements(Tensor.index_select)
+def index_select(
+    input: EdgeIndex,
+    dim: int,
+    index: Tensor,
+) -> Union[EdgeIndex, Tensor]:
+
+    out = Tensor.__torch_function__(  #
+        torch.index_select, (Tensor, ), (input, dim, index))
+
+    if dim == 1 or dim == -1:
+        out = out.as_subclass(EdgeIndex)
+        out._sparse_size = input.sparse_size
+
+    return out
+
+
+@implements(torch.narrow)
+@implements(Tensor.narrow)
+def narrow(
+    input: EdgeIndex,
+    dim: int,
+    start: Union[int, Tensor],
+    length: int,
+) -> Union[EdgeIndex, Tensor]:
+
+    out = Tensor.__torch_function__(  #
+        torch.narrow, (Tensor, ), (input, dim, start, length))
+
+    if dim == 1 or dim == -1:
+        out = out.as_subclass(EdgeIndex)
+        out._sparse_size = input.sparse_size
+        # NOTE We could potentially maintain `rowptr`/`colptr` attributes here,
+        # but it is not really clear if this is worth it. The most important
+        # information, the sort order, needs to be maintained though:
+        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 eddd10879..b62d818f2 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -4,7 +4,7 @@ from torch_geometric.data.edge_index import EdgeIndex
 from torch_geometric.testing import onlyCUDA, withCUDA
 
 
-def test_edge_index():
+def test_basic():
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sparse_size=(3, 3))
     assert isinstance(adj, EdgeIndex)
     assert str(adj) == ('EdgeIndex([[0, 1, 1, 2],\n'
@@ -17,7 +17,7 @@ def test_edge_index():
     assert not isinstance(adj + 1, EdgeIndex)
 
 
-def test_edge_index_fill_cache():
+def test_fill_cache():
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
     adj.validate().fill_cache()
     assert adj.sparse_size == (3, None)
@@ -30,7 +30,7 @@ def test_edge_index_fill_cache():
 
 
 @withCUDA
-def test_edge_index_to(device):
+def test_to(device):
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]])
 
     out = adj.to(torch.int)
@@ -47,7 +47,7 @@ def test_edge_index_to(device):
 
 
 @onlyCUDA
-def test_edge_index_cpu_cuda():
+def test_cpu_cuda():
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]])
 
     out = adj.cuda()
@@ -59,7 +59,7 @@ def test_edge_index_cpu_cuda():
     assert not out.is_cuda
 
 
-def test_edge_index_sort_by():
+def test_sort_by():
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
     out = adj.sort_by('row')
     assert isinstance(out, torch.return_types.sort)
@@ -104,7 +104,7 @@ def test_edge_index_sort_by():
     assert torch.equal(out.values._csc2csr, torch.tensor([1, 0, 3, 2]))
 
 
-def test_edge_index_cat():
+def test_cat():
     adj1 = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sparse_size=(3, 3))
     adj2 = EdgeIndex([[1, 2, 2, 3], [2, 1, 3, 2]], sparse_size=(4, 4))
 
@@ -119,7 +119,7 @@ def test_edge_index_cat():
     assert not isinstance(out, EdgeIndex)
 
 
-def test_edge_index_flip():
+def test_flip():
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
     adj.fill_cache()
 
@@ -136,3 +136,28 @@ def test_edge_index_flip():
     assert out.sparse_size == (None, 3)
     assert out.sort_order is None
     assert out._colptr is None
+
+
+def test_index_select():
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
+
+    out = adj.index_select(1, torch.tensor([1, 3]))
+    assert torch.equal(out, torch.tensor([[1, 2], [0, 1]]))
+    assert isinstance(out, EdgeIndex)
+
+    out = adj.index_select(0, torch.tensor([0]))
+    assert torch.equal(out, torch.tensor([[0, 1, 1, 2]]))
+    assert not isinstance(out, EdgeIndex)
+
+
+def test_narrow():
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
+
+    out = adj.narrow(dim=1, start=1, length=2)
+    assert torch.equal(out, torch.tensor([[1, 1], [0, 2]]))
+    assert isinstance(out, EdgeIndex)
+    assert out.sort_order == 'row'
+
+    out = adj.narrow(dim=0, start=0, length=1)
+    assert torch.equal(out, torch.tensor([[0, 1, 1, 2]]))
+    assert not isinstance(out, EdgeIndex)