Repository: pyg-team/pytorch_geometric
Pull request #7361 (2023-05-13 11:16:05)

Added `FilterEdges` graph coarsening operator

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7361?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7361](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7361?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (ddce1a7) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/6092039ebcaa2f769d8f11eb919a1cc2bcfb3db1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (6092039) will **decrease** coverage by `0.30%`.
> The diff coverage is `100.00%`.

> :exclamation: Current head ddce1a7 differs from pull request most recent head 956b68c. Consider uploading reports for the commit 956b68c to get more accurate results

```diff
@@            Coverage Diff             @@
##           master    #7361      +/-   ##
==========================================
- Coverage   91.95%   91.65%   -0.30%     
==========================================
  Files         438      439       +1     
  Lines       24367    24392      +25     
==========================================
- Hits        22406    22356      -50     
- Misses       1961     2036      +75     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7361?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/nn/pool/connect/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7361?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvY29ubmVjdC9fX2luaXRfXy5weQ==) | `100.00% <100.00%> (+100.00%)` | :arrow_up: |
| [torch\_geometric/nn/pool/connect/base.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7361?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvY29ubmVjdC9iYXNlLnB5) | `93.10% <100.00%> (+93.10%)` | :arrow_up: |
| [torch\_geometric/nn/pool/connect/filter\_edges.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7361?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvY29ubmVjdC9maWx0ZXJfZWRnZXMucHk=) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/nn/pool/topk\_pool.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7361?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvdG9wa19wb29sLnB5) | `100.00% <100.00%> (ø)` | |

... and [20 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7361/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

: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)

puririshi98: LGTM besides dropping the old filter_adj function: https://github.com/pyg-team/pytorch_geometric/pull/7361/files#diff-eeda03d73110e079fdee8218bb69b7021bf6b21158471c6d07edb9a41c40f934R11
wsad1: @puririshi98 any reason you want to keep the old `filter_adj` function. Is it used in other places?
puririshi98: > @puririshi98 any reason you want to keep the old `filter_adj` function. Is it used in other places?

no reason, in my previous comment I mentioned
`LGTM besides dropping the old filter_adj function: https://github.com/pyg-team/pytorch_geometric/pull/7361/files#diff-eeda03d73110e079fdee8218bb69b7021bf6b21158471c6d07edb9a41c40f934R11`
puririshi98: turns out its needed for pan pool: https://github.com/pyg-team/pytorch_geometric/actions/runs/5258060418/jobs/9501770183?pr=7361
will fix
wsad1: @puririshi98 feel free to approve it if you feel its ready. We can merge after that.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6fb18c718..55b609ebe 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+<<<<<<< HEAD
+- Added `FilterEdges` graph coarsening operator ([#7361](https://github.com/pyg-team/pytorch_geometric/pull/7361))
+=======
 - Added the `DirGNN` model for learning on directed graphs ([#7458](https://github.com/pyg-team/pytorch_geometric/pull/7458))
 - Allow GPU tensors as input to `NodeLoader` and `LinkLoader` ([#7572](https://github.com/pyg-team/pytorch_geometric/pull/7572))
 - Added an `embedding_device` option to allow for GPU inference in `BasicGNN` ([#7548](https://github.com/pyg-team/pytorch_geometric/pull/7548))
@@ -30,6 +33,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added padding capabilities to `HeteroData.to_homogeneous()` in case feature dimensionalities do not match ([#7374](https://github.com/pyg-team/pytorch_geometric/pull/7374))
 - Added an optional `batch_size` argument to `fps`, `knn`, `knn_graph`, `radius` and `radius_graph` ([#7368](https://github.com/pyg-team/pytorch_geometric/pull/7368))
 - Added `PrefetchLoader` capabilities ([#7376](https://github.com/pyg-team/pytorch_geometric/pull/7376), [#7378](https://github.com/pyg-team/pytorch_geometric/pull/7378), [#7383](https://github.com/pyg-team/pytorch_geometric/pull/7383))
+>>>>>>> 8444e4869e20d45f5ba27f5dab8c455b3fa36094
 - Added an example for hierarichial sampling ([#7244](https://github.com/pyg-team/pytorch_geometric/pull/7244))
 - Added Kùzu remote backend examples ([#7298](https://github.com/pyg-team/pytorch_geometric/pull/7298))
 - Fixed tracing of `add_self_loops` for a dynamic number of nodes ([#7330](https://github.com/pyg-team/pytorch_geometric/pull/7330))
diff --git a/torch_geometric/nn/pool/connect/__init__.py b/torch_geometric/nn/pool/connect/__init__.py
index 52b45639f..1e24f3ee0 100644
--- a/torch_geometric/nn/pool/connect/__init__.py
+++ b/torch_geometric/nn/pool/connect/__init__.py
@@ -1,6 +1,8 @@
 from .base import Connect, ConnectOutput
+from .filter_edges import FilterEdges
 
 __all__ = [
     'Connect',
     'ConnectOutput',
+    'FilterEdges',
 ]
diff --git a/torch_geometric/nn/pool/connect/base.py b/torch_geometric/nn/pool/connect/base.py
index fd78a5c61..23d0dd443 100644
--- a/torch_geometric/nn/pool/connect/base.py
+++ b/torch_geometric/nn/pool/connect/base.py
@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Optional, Tuple
+from typing import Optional
 
 import torch
 from torch import Tensor
@@ -70,7 +70,7 @@ class Connect(torch.nn.Module):
         edge_index: Tensor,
         edge_attr: Optional[Tensor] = None,
         batch: Optional[Tensor] = None,
-    ) -> Tuple[Tensor, Optional[Tensor]]:
+    ) -> ConnectOutput:
         r"""
         Args:
             select_output (SelectOutput): The output of :class:`Select`.
diff --git a/torch_geometric/nn/pool/connect/filter_edges.py b/torch_geometric/nn/pool/connect/filter_edges.py
new file mode 100644
index 000000000..2481f4514
--- /dev/null
+++ b/torch_geometric/nn/pool/connect/filter_edges.py
@@ -0,0 +1,61 @@
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from torch_geometric.nn.pool.select import SelectOutput
+from torch_geometric.utils.num_nodes import maybe_num_nodes
+
+from .base import Connect, ConnectOutput
+
+
+def filter_adj(
+    edge_index: Tensor,
+    edge_attr: Optional[Tensor],
+    node_index: Tensor,
+    cluster_index: Optional[Tensor] = None,
+    num_nodes: Optional[int] = None,
+) -> Tuple[Tensor, Optional[Tensor]]:
+    num_nodes = maybe_num_nodes(edge_index, num_nodes)
+    if cluster_index is None:
+        cluster_index = torch.arange(node_index.size(0), dtype=torch.long,
+                                     device=node_index.device)
+    mask = node_index.new_full((num_nodes, ), -1)
+    mask[node_index] = cluster_index
+
+    row, col = edge_index[0], edge_index[1]
+    row, col = mask[row], mask[col]
+    mask = (row >= 0) & (col >= 0)
+    row, col = row[mask], col[mask]
+
+    if edge_attr is not None:
+        edge_attr = edge_attr[mask]
+
+    return torch.stack([row, col], dim=0), edge_attr
+
+
+class FilterEdges(Connect):
+    r"""Filter out edges if their incident nodes are not in any cluster.
+    It is assumed that each cluster contains only one node.
+
+    .. math::
+            \mathbf{A}^{\prime} &= \mathbf{A}_{\mathbf{i},\mathbf{i}}
+
+    Where :math:`\mathbf{i}` are retained nodes.
+    """
+    def forward(
+        self,
+        select_output: SelectOutput,
+        edge_index: Tensor,
+        edge_attr: Optional[Tensor] = None,
+        batch: Optional[Tensor] = None,
+    ) -> ConnectOutput:
+        if select_output.num_clusters != select_output.cluster_index.size(0):
+            raise ValueError("'FilterEdges' requires each cluster to contain "
+                             "only one node.")
+        edge_index, edge_attr = filter_adj(edge_index, edge_attr,
+                                           select_output.node_index,
+                                           select_output.cluster_index,
+                                           num_nodes=select_output.num_nodes)
+        batch = self.get_pooled_batch(select_output, batch)
+        return ConnectOutput(edge_index, edge_attr, batch)
diff --git a/torch_geometric/nn/pool/pan_pool.py b/torch_geometric/nn/pool/pan_pool.py
index f0c0ab66d..763a6148d 100644
--- a/torch_geometric/nn/pool/pan_pool.py
+++ b/torch_geometric/nn/pool/pan_pool.py
@@ -4,8 +4,8 @@ import torch
 from torch import Tensor
 from torch.nn import Parameter
 
+from torch_geometric.nn.pool.connect.filter_edges import filter_adj
 from torch_geometric.nn.pool.select.topk import topk
-from torch_geometric.nn.pool.topk_pool import filter_adj
 from torch_geometric.typing import OptTensor, SparseTensor
 from torch_geometric.utils import scatter, softmax
 
@@ -101,7 +101,8 @@ class PANPooling(torch.nn.Module):
         x = self.multiplier * x if self.multiplier != 1 else x
 
         edge_index = torch.stack([col, row], dim=0)
-        edge_index, edge_weight = filter_adj(edge_index, edge_weight, perm,
+        edge_index, edge_weight = filter_adj(edge_index, edge_weight,
+                                             node_index=perm,
                                              num_nodes=score.size(0))
         assert edge_weight is not None
 
diff --git a/torch_geometric/nn/pool/sag_pool.py b/torch_geometric/nn/pool/sag_pool.py
index 580467949..d205288a2 100644
--- a/torch_geometric/nn/pool/sag_pool.py
+++ b/torch_geometric/nn/pool/sag_pool.py
@@ -4,8 +4,8 @@ import torch
 from torch import Tensor
 
 from torch_geometric.nn import GraphConv
+from torch_geometric.nn.pool.connect.filter_edges import filter_adj
 from torch_geometric.nn.pool.select.topk import topk
-from torch_geometric.nn.pool.topk_pool import filter_adj
 from torch_geometric.typing import OptTensor
 from torch_geometric.utils import softmax
 
@@ -136,8 +136,9 @@ class SAGPooling(torch.nn.Module):
         x = self.multiplier * x if self.multiplier != 1 else x
 
         batch = batch[perm]
-        edge_index, edge_attr = filter_adj(edge_index, edge_attr, perm,
-                                           num_nodes=score.size(0))
+        edge_index, edge_weight = filter_adj(edge_index, edge_attr,
+                                             node_index=perm,
+                                             num_nodes=score.size(0))
 
         return x, edge_index, edge_attr, batch, perm, score[perm]
 
diff --git a/torch_geometric/nn/pool/topk_pool.py b/torch_geometric/nn/pool/topk_pool.py
index 5560c9d3a..dc6d65f70 100644
--- a/torch_geometric/nn/pool/topk_pool.py
+++ b/torch_geometric/nn/pool/topk_pool.py
@@ -3,31 +3,8 @@ from typing import Callable, Optional, Tuple, Union
 import torch
 from torch import Tensor
 
+from torch_geometric.nn.pool.connect import FilterEdges
 from torch_geometric.nn.pool.select import SelectTopK
-from torch_geometric.utils.num_nodes import maybe_num_nodes
-
-
-def filter_adj(
-    edge_index: Tensor,
-    edge_attr: Optional[Tensor],
-    perm: Tensor,
-    num_nodes: Optional[int] = None,
-) -> Tuple[Tensor, Optional[Tensor]]:
-    num_nodes = maybe_num_nodes(edge_index, num_nodes)
-
-    mask = perm.new_full((num_nodes, ), -1)
-    i = torch.arange(perm.size(0), dtype=torch.long, device=perm.device)
-    mask[perm] = i
-
-    row, col = edge_index[0], edge_index[1]
-    row, col = mask[row], mask[col]
-    mask = (row >= 0) & (col >= 0)
-    row, col = row[mask], col[mask]
-
-    if edge_attr is not None:
-        edge_attr = edge_attr[mask]
-
-    return torch.stack([row, col], dim=0), edge_attr
 
 
 class TopKPooling(torch.nn.Module):
@@ -100,6 +77,7 @@ class TopKPooling(torch.nn.Module):
         self.multiplier = multiplier
 
         self.select = SelectTopK(in_channels, ratio, min_score, nonlinearity)
+        self.connect = FilterEdges()
 
         self.reset_parameters()
 
@@ -141,9 +119,11 @@ class TopKPooling(torch.nn.Module):
         x = x[perm] * score.view(-1, 1)
         x = self.multiplier * x if self.multiplier != 1 else x
 
-        batch = batch[perm]
-        edge_index, edge_attr = filter_adj(edge_index, edge_attr, perm,
-                                           num_nodes=select_output.num_nodes)
+        connect_output = self.connect(select_output, edge_index, edge_attr,
+                                      batch)
+        edge_index = connect_output.edge_index
+        edge_attr = connect_output.edge_attr
+        batch = connect_output.batch
 
         return x, edge_index, edge_attr, batch, perm, score
 


Test patch:
diff --git a/test/nn/pool/connect/test_filter_edges.py b/test/nn/pool/connect/test_filter_edges.py
new file mode 100644
index 000000000..4e1f840bf
--- /dev/null
+++ b/test/nn/pool/connect/test_filter_edges.py
@@ -0,0 +1,28 @@
+import torch
+
+from torch_geometric.nn.pool.connect import FilterEdges
+from torch_geometric.nn.pool.select import SelectOutput
+from torch_geometric.testing import is_full_test
+
+
+def test_filter_edges():
+    edge_index = torch.tensor([[0, 1, 1, 2, 2, 3], [1, 0, 1, 3, 2, 2]])
+    edge_attr = torch.tensor([1, 2, 3, 4, 5, 6])
+    batch = torch.tensor([0, 0, 1, 1])
+
+    node_index = torch.tensor([1, 2])
+    cluster_index = torch.tensor([0, 1])
+    select_output = SelectOutput(node_index=node_index, num_nodes=4,
+                                 cluster_index=cluster_index, num_clusters=2)
+    connect = FilterEdges()
+    output1 = connect(select_output, edge_index, edge_attr, batch)
+    assert output1.edge_index.tolist() == [[0, 1], [0, 1]]
+    assert output1.edge_attr.tolist() == [3, 5]
+    assert output1.batch.tolist() == [0, 1]
+
+    if is_full_test():
+        jit = torch.jit.script(connect)
+        output2 = jit(select_output, edge_index, edge_attr, batch)
+        torch.allclose(output1.edge_index, output2.edge_index)
+        torch.allclose(output1.edge_attr, output2.edge_attr)
+        torch.allclose(output1.batch, output2.batch)
diff --git a/test/nn/pool/test_topk_pool.py b/test/nn/pool/test_topk_pool.py
index a0616183f..5e3fb2095 100644
--- a/test/nn/pool/test_topk_pool.py
+++ b/test/nn/pool/test_topk_pool.py
@@ -1,7 +1,7 @@
 import torch
 
 from torch_geometric.nn.pool import TopKPooling
-from torch_geometric.nn.pool.topk_pool import filter_adj
+from torch_geometric.nn.pool.connect.filter_edges import filter_adj
 from torch_geometric.testing import is_full_test