Repository: pyg-team/pytorch_geometric
Pull request #6847 (2023-03-03 13:20:05)

Added PyTorch Sparse Tensor support for `remove_self_loops`, `add_self_loops`, and `maybe_num_nodes`
This PR aims to add PyTorch Sparse Tensor support for `remove_self_loops`, `add_self_loops`, and `maybe_num_nodes`, which will be used for supporting PyTorch Sparse Tensor in `torch_geometric.nn.conv`, e.g. `AGNNConv`:
https://github.com/pyg-team/pytorch_geometric/blob/09dbb20176b3de850317277329c63d756504c783/torch_geometric/nn/conv/agnn_conv.py#L65-L69

There are some questions to be clarified:
+ Should I also add support for other functions like `add_remaining_self_loops`?
+ Should I also update the doc-strings correspondingly?


codecov[bot]: # [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6847?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#6847](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6847?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (ba36ae5) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/09dbb20176b3de850317277329c63d756504c783?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (09dbb20) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #6847   +/-   ##
=======================================
  Coverage   91.49%   91.50%           
=======================================
  Files         427      427           
  Lines       23057    23076   +19     
=======================================
+ Hits        21097    21116   +19     
  Misses       1960     1960           
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6847?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/utils/loop.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6847?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL2xvb3AucHk=) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/utils/num\_nodes.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/6847?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL251bV9ub2Rlcy5weQ==) | `96.77% <100.00%> (+0.34%)` | :arrow_up: |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e884c44f9..d095ef2b8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -91,7 +91,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `Data.edge_subgraph` and `HeteroData.edge_subgraph` functionalities ([#6193](https://github.com/pyg-team/pytorch_geometric/pull/6193))
 - Added `input_time` option to `LightningNodeData` and `transform_sampler_output` to `NodeLoader` and `LinkLoader` ([#6187](https://github.com/pyg-team/pytorch_geometric/pull/6187))
 - Added `summary` for PyG/PyTorch models ([#5859](https://github.com/pyg-team/pytorch_geometric/pull/5859), [#6161](https://github.com/pyg-team/pytorch_geometric/pull/6161))
-- Started adding `torch.sparse` support to PyG ([#5906](https://github.com/pyg-team/pytorch_geometric/pull/5906), [#5944](https://github.com/pyg-team/pytorch_geometric/pull/5944), [#6003](https://github.com/pyg-team/pytorch_geometric/pull/6003), [#6033](https://github.com/pyg-team/pytorch_geometric/pull/6033), [#6514](https://github.com/pyg-team/pytorch_geometric/pull/6514), [#6532](https://github.com/pyg-team/pytorch_geometric/pull/6532), [#6748](https://github.com/pyg-team/pytorch_geometric/pull/6748))
+- Started adding `torch.sparse` support to PyG ([#5906](https://github.com/pyg-team/pytorch_geometric/pull/5906), [#5944](https://github.com/pyg-team/pytorch_geometric/pull/5944), [#6003](https://github.com/pyg-team/pytorch_geometric/pull/6003), [#6033](https://github.com/pyg-team/pytorch_geometric/pull/6033), [#6514](https://github.com/pyg-team/pytorch_geometric/pull/6514), [#6532](https://github.com/pyg-team/pytorch_geometric/pull/6532), [#6748](https://github.com/pyg-team/pytorch_geometric/pull/6748), [#6847](https://github.com/pyg-team/pytorch_geometric/pull/6847))
 - Add `inputs_channels` back in training benchmark ([#6154](https://github.com/pyg-team/pytorch_geometric/pull/6154))
 - Added support for dropping nodes in `utils.to_dense_batch` in case `max_num_nodes` is smaller than the number of nodes ([#6124](https://github.com/pyg-team/pytorch_geometric/pull/6124))
 - Added the RandLA-Net architecture as an example ([#5117](https://github.com/pyg-team/pytorch_geometric/pull/5117))
diff --git a/torch_geometric/utils/loop.py b/torch_geometric/utils/loop.py
index 91763f9bb..4a588c66b 100644
--- a/torch_geometric/utils/loop.py
+++ b/torch_geometric/utils/loop.py
@@ -7,6 +7,7 @@ from torch_geometric.typing import OptTensor
 from torch_geometric.utils import scatter
 
 from .num_nodes import maybe_num_nodes
+from .sparse import is_torch_sparse_tensor, to_edge_index, to_torch_coo_tensor
 
 
 def contains_self_loops(edge_index: Tensor) -> bool:
@@ -60,8 +61,20 @@ def remove_self_loops(
         tensor([[1, 2],
                 [3, 4]]))
     """
+    is_sparse = is_torch_sparse_tensor(edge_index)
+    if is_sparse:
+        assert edge_attr is None
+        size = edge_index.size()
+        edge_index, edge_attr = to_edge_index(edge_index)
+
     mask = edge_index[0] != edge_index[1]
     edge_index = edge_index[:, mask]
+
+    if is_sparse:
+        edge_attr = edge_attr[mask]
+        adj = to_torch_coo_tensor(edge_index, edge_attr, size=size)
+        return adj, None
+
     if edge_attr is None:
         return edge_index, None
     else:
@@ -185,6 +198,11 @@ def add_self_loops(
     """
     N = maybe_num_nodes(edge_index, num_nodes)
 
+    is_sparse = is_torch_sparse_tensor(edge_index)
+    if is_sparse:
+        assert edge_attr is None
+        edge_index, edge_attr = to_edge_index(edge_index)
+
     loop_index = torch.arange(0, N, dtype=torch.long, device=edge_index.device)
     loop_index = loop_index.unsqueeze(0).repeat(2, 1)
 
@@ -211,6 +229,8 @@ def add_self_loops(
         edge_attr = torch.cat([edge_attr, loop_attr], dim=0)
 
     edge_index = torch.cat([edge_index, loop_index], dim=1)
+    if is_sparse:
+        return to_torch_coo_tensor(edge_index, edge_attr), None
     return edge_index, edge_attr
 
 
diff --git a/torch_geometric/utils/num_nodes.py b/torch_geometric/utils/num_nodes.py
index abf3acb90..0bbe04561 100644
--- a/torch_geometric/utils/num_nodes.py
+++ b/torch_geometric/utils/num_nodes.py
@@ -6,6 +6,8 @@ from torch import Tensor
 
 from torch_geometric.typing import SparseTensor  # noqa
 
+from .sparse import is_torch_sparse_tensor
+
 
 @torch.jit._overload
 def maybe_num_nodes(edge_index, num_nodes=None):
@@ -23,6 +25,8 @@ def maybe_num_nodes(edge_index, num_nodes=None):
     if num_nodes is not None:
         return num_nodes
     elif isinstance(edge_index, Tensor):
+        if is_torch_sparse_tensor(edge_index):
+            return max(edge_index.size(0), edge_index.size(1))
         return int(edge_index.max()) + 1 if edge_index.numel() > 0 else 0
     else:
         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 f745c70d0..9e54a57a4 100644
--- a/test/utils/test_loop.py
+++ b/test/utils/test_loop.py
@@ -7,6 +7,7 @@ from torch_geometric.utils import (
     get_self_loop_attr,
     remove_self_loops,
     segregate_self_loops,
+    to_torch_coo_tensor,
 )
 
 
@@ -32,6 +33,10 @@ def test_remove_self_loops():
     assert out[0].tolist() == expected
     assert out[1].tolist() == [[1, 2], [3, 4]]
 
+    adj = to_torch_coo_tensor(edge_index)
+    adj, _ = remove_self_loops(adj)
+    assert torch.diag(adj.to_dense()).tolist() == [0, 0]
+
 
 def test_segregate_self_loops():
     edge_index = torch.tensor([[0, 0, 1], [0, 1, 0]])
@@ -54,6 +59,7 @@ def test_add_self_loops():
     edge_index = torch.tensor([[0, 1, 0], [1, 0, 0]])
     edge_weight = torch.tensor([0.5, 0.5, 0.5])
     edge_attr = torch.eye(3)
+    adj = to_torch_coo_tensor(edge_index, edge_weight)
 
     expected = [[0, 1, 0, 0, 1], [1, 0, 0, 0, 1]]
     assert add_self_loops(edge_index)[0].tolist() == expected
@@ -62,14 +68,26 @@ def test_add_self_loops():
     assert out[0].tolist() == expected
     assert out[1].tolist() == [0.5, 0.5, 0.5, 1., 1.]
 
+    out = add_self_loops(adj)[0]
+    assert out._indices().tolist() == [[0, 0, 1, 1], [0, 1, 0, 1]]
+    assert out._values().tolist() == [1.5, 0.5, 0.5, 1.0]
+
     out = add_self_loops(edge_index, edge_weight, fill_value=5)
     assert out[0].tolist() == expected
     assert out[1].tolist() == [0.5, 0.5, 0.5, 5.0, 5.0]
 
+    out = add_self_loops(adj, fill_value=5)[0]
+    assert out._indices().tolist() == [[0, 0, 1, 1], [0, 1, 0, 1]]
+    assert out._values().tolist() == [5.5, 0.5, 0.5, 5.0]
+
     out = add_self_loops(edge_index, edge_weight, fill_value=torch.tensor(2.))
     assert out[0].tolist() == expected
     assert out[1].tolist() == [0.5, 0.5, 0.5, 2., 2.]
 
+    out = add_self_loops(adj, fill_value=torch.tensor(2.))[0]
+    assert out._indices().tolist() == [[0, 0, 1, 1], [0, 1, 0, 1]]
+    assert out._values().tolist() == [2.5, 0.5, 0.5, 2.0]
+
     out = add_self_loops(edge_index, edge_weight, fill_value='add')
     assert out[0].tolist() == expected
     assert out[1].tolist() == [0.5, 0.5, 0.5, 1, 0.5]
diff --git a/test/utils/test_num_nodes.py b/test/utils/test_num_nodes.py
index d1723b149..ce18d6113 100644
--- a/test/utils/test_num_nodes.py
+++ b/test/utils/test_num_nodes.py
@@ -1,5 +1,6 @@
 import torch
 
+from torch_geometric.utils import to_torch_coo_tensor
 from torch_geometric.utils.num_nodes import (
     maybe_num_nodes,
     maybe_num_nodes_dict,
@@ -12,6 +13,10 @@ def test_maybe_num_nodes():
     assert maybe_num_nodes(edge_index, 4) == 4
     assert maybe_num_nodes(edge_index) == 3
 
+    adj = to_torch_coo_tensor(edge_index)
+    assert maybe_num_nodes(adj, 4) == 4
+    assert maybe_num_nodes(adj) == 3
+
 
 def test_maybe_num_nodes_dict():
     edge_index_dict = {