Repository: pyg-team/pytorch_geometric
Pull request #8655 (2023-12-21 10:13:21)

Support `EdgeIndex` in `remove_self_loops` and `segregate_self_loops`


Patch:
diff --git a/torch_geometric/utils/loop.py b/torch_geometric/utils/loop.py
index c76c78c40..4ddebe95b 100644
--- a/torch_geometric/utils/loop.py
+++ b/torch_geometric/utils/loop.py
@@ -182,11 +182,20 @@ def segregate_self_loops(  # noqa: F811
     mask = edge_index[0] != edge_index[1]
     inv_mask = ~mask
 
+    is_undirected = False
+    if not torch.jit.is_scripting() and isinstance(edge_index, EdgeIndex):
+        is_undirected = edge_index.is_undirected
+
     loop_edge_index = edge_index[:, inv_mask]
     loop_edge_attr = None if edge_attr is None else edge_attr[inv_mask]
     edge_index = edge_index[:, mask]
     edge_attr = None if edge_attr is None else edge_attr[mask]
 
+    if not torch.jit.is_scripting() and isinstance(edge_index, EdgeIndex):
+        assert isinstance(loop_edge_index, EdgeIndex)
+        edge_index._is_undirected = is_undirected
+        loop_edge_index._is_undirected = is_undirected
+
     return edge_index, edge_attr, loop_edge_index, loop_edge_attr
 
 


Test patch:
diff --git a/test/utils/test_loop.py b/test/utils/test_loop.py
index 9287f9079..f9f09ec0a 100644
--- a/test/utils/test_loop.py
+++ b/test/utils/test_loop.py
@@ -1,5 +1,6 @@
 import torch
 
+from torch_geometric import EdgeIndex
 from torch_geometric.utils import (
     add_remaining_self_loops,
     add_self_loops,
@@ -37,6 +38,26 @@ def test_remove_self_loops():
     adj, _ = remove_self_loops(adj)
     assert torch.diag(adj.to_dense()).tolist() == [0, 0]
 
+    edge_index = EdgeIndex(
+        edge_index,
+        sparse_size=(2, 2),
+        sort_order='row',
+        is_undirected=True,
+    )
+    out = remove_self_loops(edge_index)
+    assert out[0].tolist() == expected
+    assert out[0].sparse_size() == (2, 2)
+    assert out[0].sort_order == 'row'
+    assert out[0].is_undirected
+    assert out[1] is None
+
+    out = remove_self_loops(edge_index, edge_attr)
+    assert out[0].tolist() == expected
+    assert out[0].sparse_size() == (2, 2)
+    assert out[0].sort_order == 'row'
+    assert out[0].is_undirected
+    assert out[1].tolist() == [[1, 2], [3, 4]]
+
 
 def test_segregate_self_loops():
     edge_index = torch.tensor([[0, 0, 1], [0, 1, 0]])
@@ -54,6 +75,36 @@ def test_segregate_self_loops():
     assert out[2].tolist() == [[0], [0]]
     assert out[3].tolist() == [1]
 
+    edge_index = EdgeIndex(
+        edge_index,
+        sparse_size=(2, 2),
+        sort_order='row',
+        is_undirected=True,
+    )
+    out = segregate_self_loops(edge_index)
+    assert out[0].tolist() == [[0, 1], [1, 0]]
+    assert out[0].sparse_size() == (2, 2)
+    assert out[0].sort_order == 'row'
+    assert out[0].is_undirected
+    assert out[1] is None
+    assert out[2].tolist() == [[0], [0]]
+    assert out[2].sparse_size() == (2, 2)
+    assert out[2].sort_order == 'row'
+    assert out[2].is_undirected
+    assert out[3] is None
+
+    out = segregate_self_loops(edge_index, edge_attr)
+    assert out[0].tolist() == [[0, 1], [1, 0]]
+    assert out[0].sparse_size() == (2, 2)
+    assert out[0].sort_order == 'row'
+    assert out[0].is_undirected
+    assert out[1].tolist() == [2, 3]
+    assert out[2].tolist() == [[0], [0]]
+    assert out[2].sparse_size() == (2, 2)
+    assert out[2].sort_order == 'row'
+    assert out[2].is_undirected
+    assert out[3].tolist() == [1]
+
 
 def test_add_self_loops():
     edge_index = torch.tensor([[0, 1, 0], [1, 0, 0]])