Repository: pyg-team/pytorch_geometric
Pull request #8665 (2023-12-22 13:32:40)

Support `EdgeIndex` in `add_self_loops`


Patch:
diff --git a/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
index c38fc00cf..67653db20 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -1008,6 +1008,9 @@ def cat(
     if dim != 1 and dim != -1:  # No valid `EdgeIndex` anymore.
         return output
 
+    if any([not isinstance(tensor, EdgeIndex) for tensor in tensors]):
+        return output
+
     output = output.as_subclass(EdgeIndex)
 
     # Post-process `sparse_size`:
@@ -1030,14 +1033,10 @@ def cat(
     output._sparse_size = (num_rows, num_cols)
 
     # Post-process `is_undirected`:
-    is_undirected = True
-    for tensor in tensors:
-        if isinstance(tensor, EdgeIndex):
-            is_undirected = tensor.is_undirected
-        else:
-            is_undirected = False
-
-    output._is_undirected = is_undirected
+    output._is_undirected = all([
+        isinstance(tensor, EdgeIndex) and tensor.is_undirected
+        for tensor in tensors
+    ])
 
     return output
 
diff --git a/torch_geometric/utils/loop.py b/torch_geometric/utils/loop.py
index 4ddebe95b..c4e5e0de7 100644
--- a/torch_geometric/utils/loop.py
+++ b/torch_geometric/utils/loop.py
@@ -459,8 +459,15 @@ def add_self_loops(  # noqa: F811
         N = maybe_num_nodes(edge_index, num_nodes)
         size = (N, N)
 
-    loop_index = torch.arange(0, N, device=edge_index.device)
-    loop_index = loop_index.unsqueeze(0).repeat(2, 1)
+    device = edge_index.device
+    if torch.jit.is_scripting():
+        loop_index = torch.arange(0, N, device=device).view(1, -1).repeat(2, 1)
+    else:
+        loop_index = EdgeIndex(
+            torch.arange(0, N, device=device).view(1, -1).repeat(2, 1),
+            sparse_size=(N, N),
+            is_undirected=True,
+        )
     full_edge_index = torch.cat([edge_index, loop_index], dim=1)
 
     if is_sparse:
diff --git a/torch_geometric/utils/num_nodes.py b/torch_geometric/utils/num_nodes.py
index 1bc337e7d..7142a045c 100644
--- a/torch_geometric/utils/num_nodes.py
+++ b/torch_geometric/utils/num_nodes.py
@@ -6,6 +6,7 @@ import torch
 from torch import Tensor
 
 import torch_geometric
+from torch_geometric import EdgeIndex
 from torch_geometric.typing import EdgeType, NodeType, SparseTensor
 
 if typing.TYPE_CHECKING:
@@ -44,6 +45,8 @@ def maybe_num_nodes(  # noqa: F811
 ) -> int:
     if num_nodes is not None:
         return num_nodes
+    elif not torch.jit.is_scripting() and isinstance(edge_index, EdgeIndex):
+        return max(edge_index.get_sparse_size())
     elif isinstance(edge_index, Tensor):
         if torch_geometric.utils.is_torch_sparse_tensor(edge_index):
             return max(edge_index.size(0), edge_index.size(1))


Test patch:
diff --git a/test/utils/test_loop.py b/test/utils/test_loop.py
index f9f09ec0a..a381bbb0f 100644
--- a/test/utils/test_loop.py
+++ b/test/utils/test_loop.py
@@ -160,6 +160,18 @@ def test_add_self_loops():
     assert out[1].tolist() == [[1., 0., 0.], [0., 1., 0.], [0., 0., 1.],
                                [0., 1., 1.], [1., 0., 0.]]
 
+    edge_index = EdgeIndex(
+        edge_index,
+        sparse_size=(2, 2),
+        sort_order='row',
+        is_undirected=True,
+    )
+    out, _ = add_self_loops(edge_index)
+    assert out.tolist() == expected
+    assert out.sparse_size() == (2, 2)
+    assert out.sort_order is None
+    assert out.is_undirected
+
     # Test empty `edge_index` and `edge_weight`:
     edge_index = torch.empty(2, 0, dtype=torch.long)
     edge_weight = torch.empty(0)