Repository: pyg-team/pytorch_geometric
Pull request #7308 (2023-05-06 06:00:18)

Refactor `TopkPooling` into `SelectTopK`
Towards #6455.
Do not review before #7307 is merged.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7308?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7308](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7308?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (140ef22) into [pool_update](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/93d4d307b094044b0fa1d37a5561c300f47b9349?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (93d4d30) will **increase** coverage by `0.12%`.
> The diff coverage is `100.00%`.

```diff
@@               Coverage Diff               @@
##           pool_update    #7308      +/-   ##
===============================================
+ Coverage        91.52%   91.64%   +0.12%     
===============================================
  Files              437      439       +2     
  Lines            24271    24312      +41     
===============================================
+ Hits             22214    22281      +67     
+ Misses            2057     2031      -26     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7308?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/asap.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7308?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvYXNhcC5weQ==) | `90.41% <100.00%> (ø)` | |
| [torch\_geometric/nn/pool/connect/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7308?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/topk\_connect.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7308?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvY29ubmVjdC90b3BrX2Nvbm5lY3QucHk=) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/nn/pool/pan\_pool.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7308?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvcGFuX3Bvb2wucHk=) | `95.74% <100.00%> (+0.09%)` | :arrow_up: |
| [torch\_geometric/nn/pool/sag\_pool.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7308?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvc2FnX3Bvb2wucHk=) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/nn/pool/select/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7308?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvc2VsZWN0L19faW5pdF9fLnB5) | `100.00% <100.00%> (+100.00%)` | :arrow_up: |
| [torch\_geometric/nn/pool/select/base.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7308?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvc2VsZWN0L2Jhc2UucHk=) | `93.75% <100.00%> (+93.75%)` | :arrow_up: |
| [torch\_geometric/nn/pool/select/topk\_select.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7308?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL3Bvb2wvc2VsZWN0L3RvcGtfc2VsZWN0LnB5) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/nn/pool/topk\_pool.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7308?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 [1 file with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7308/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)

rusty1s: Hi @wsad1, I modified this PR quite a bit. Hope the changes work for you:
* I moved attention score computation to `SelectTopK` - I think this belongs there.
* Added tests for `SelectTopK`
* I reverted the addition of `TopKConnect` - Can we add that separately with corresponding test cases? I would also be in favor of renaming that to something more general, e.g., `FilterGraph`? The name should make it clear that we only maintain the nodes present in `select_output.node_index`. It may also require changes since it should respect the assignment matrix coming from `SelectOutput`
* Can you add documentation of `SelectTopK` in a follow-up?

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 73052b80b..b1b616f21 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Unify graph pooling framework ([#7308](https://github.com/pyg-team/pytorch_geometric/pull/7308))
 - Added support for tuples as keys in `ModuleDict`/`ParameterDict` ([#7294](https://github.com/pyg-team/pytorch_geometric/pull/7294))
 - Added `NodePropertySplit` transform for creating node-level splits using structural node properties ([#6894](https://github.com/pyg-team/pytorch_geometric/pull/6894))
 - Added an option to preserve directed graphs in `CitationFull` datasets ([#7275](https://github.com/pyg-team/pytorch_geometric/pull/7275))
diff --git a/torch_geometric/nn/pool/asap.py b/torch_geometric/nn/pool/asap.py
index 4fa07277a..ec391c436 100644
--- a/torch_geometric/nn/pool/asap.py
+++ b/torch_geometric/nn/pool/asap.py
@@ -7,7 +7,7 @@ from torch import Tensor
 from torch.nn import Linear
 
 from torch_geometric.nn import LEConv
-from torch_geometric.nn.pool.topk_pool import topk
+from torch_geometric.nn.pool.select.topk import topk
 from torch_geometric.utils import (
     add_remaining_self_loops,
     remove_self_loops,
diff --git a/torch_geometric/nn/pool/pan_pool.py b/torch_geometric/nn/pool/pan_pool.py
index f2f9f43a8..f0c0ab66d 100644
--- a/torch_geometric/nn/pool/pan_pool.py
+++ b/torch_geometric/nn/pool/pan_pool.py
@@ -4,7 +4,8 @@ import torch
 from torch import Tensor
 from torch.nn import Parameter
 
-from torch_geometric.nn.pool.topk_pool import filter_adj, topk
+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
 
diff --git a/torch_geometric/nn/pool/sag_pool.py b/torch_geometric/nn/pool/sag_pool.py
index 4bf80b50a..580467949 100644
--- a/torch_geometric/nn/pool/sag_pool.py
+++ b/torch_geometric/nn/pool/sag_pool.py
@@ -4,7 +4,8 @@ import torch
 from torch import Tensor
 
 from torch_geometric.nn import GraphConv
-from torch_geometric.nn.pool.topk_pool import filter_adj, topk
+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
 
diff --git a/torch_geometric/nn/pool/select/__init__.py b/torch_geometric/nn/pool/select/__init__.py
index 2a7278a33..2a724b313 100644
--- a/torch_geometric/nn/pool/select/__init__.py
+++ b/torch_geometric/nn/pool/select/__init__.py
@@ -1,6 +1,8 @@
 from .base import Select, SelectOutput
+from .topk import SelectTopK
 
 __all__ = [
     'Select',
     'SelectOutput',
+    'SelectTopK',
 ]
diff --git a/torch_geometric/nn/pool/select/base.py b/torch_geometric/nn/pool/select/base.py
index ee40f17f0..ad729776d 100644
--- a/torch_geometric/nn/pool/select/base.py
+++ b/torch_geometric/nn/pool/select/base.py
@@ -14,6 +14,7 @@ class SelectOutput(CastMixin):
 
     Args:
         node_index (torch.Tensor): The indices of the selected nodes.
+        num_nodes (int): The number of nodes.
         cluster_index (torch.Tensor): The indices of the clusters each node in
             :obj:`node_index` is assigned to.
         num_clusters (int): The number of clusters.
@@ -21,6 +22,7 @@ class SelectOutput(CastMixin):
             of the assignment of a node to its cluster. (default: :obj:`None`)
     """
     node_index: Tensor
+    num_nodes: int
     cluster_index: Tensor
     num_clusters: int
     weight: Optional[Tensor] = None
@@ -28,6 +30,7 @@ class SelectOutput(CastMixin):
     def __init__(
         self,
         node_index: Tensor,
+        num_nodes: int,
         cluster_index: Tensor,
         num_clusters: int,
         weight: Optional[Tensor] = None,
@@ -55,6 +58,7 @@ class SelectOutput(CastMixin):
                              f"values (got {weight.numel()} values)")
 
         self.node_index = node_index
+        self.num_nodes = num_nodes
         self.cluster_index = cluster_index
         self.num_clusters = num_clusters
         self.weight = weight
diff --git a/torch_geometric/nn/pool/select/topk.py b/torch_geometric/nn/pool/select/topk.py
new file mode 100644
index 000000000..aabb4a691
--- /dev/null
+++ b/torch_geometric/nn/pool/select/topk.py
@@ -0,0 +1,135 @@
+from typing import Callable, Optional, Union
+
+import torch
+from torch import Tensor
+
+from torch_geometric.nn.inits import uniform
+from torch_geometric.nn.resolver import activation_resolver
+from torch_geometric.utils import scatter, softmax
+
+from .base import Select, SelectOutput
+
+
+# TODO (matthias) Benchmark and document this method.
+def topk(
+    x: Tensor,
+    ratio: Optional[Union[float, int]],
+    batch: Tensor,
+    min_score: Optional[float] = None,
+    tol: float = 1e-7,
+) -> Tensor:
+    if min_score is not None:
+        # Make sure that we do not drop all nodes in a graph.
+        scores_max = scatter(x, batch, reduce='max')[batch] - tol
+        scores_min = scores_max.clamp(max=min_score)
+
+        perm = (x > scores_min).nonzero().view(-1)
+        return perm
+
+    if ratio is not None:
+        num_nodes = scatter(batch.new_ones(x.size(0)), batch, reduce='sum')
+        batch_size, max_num_nodes = num_nodes.size(0), int(num_nodes.max())
+
+        cum_num_nodes = torch.cat(
+            [num_nodes.new_zeros(1),
+             num_nodes.cumsum(dim=0)[:-1]], dim=0)
+
+        index = torch.arange(batch.size(0), dtype=torch.long, device=x.device)
+        index = (index - cum_num_nodes[batch]) + (batch * max_num_nodes)
+
+        dense_x = x.new_full((batch_size * max_num_nodes, ), -60000.0)
+        dense_x[index] = x
+        dense_x = dense_x.view(batch_size, max_num_nodes)
+
+        _, perm = dense_x.sort(dim=-1, descending=True)
+
+        perm = perm + cum_num_nodes.view(-1, 1)
+        perm = perm.view(-1)
+
+        if ratio >= 1:
+            k = num_nodes.new_full((num_nodes.size(0), ), int(ratio))
+            k = torch.min(k, num_nodes)
+        else:
+            k = (float(ratio) * num_nodes.to(x.dtype)).ceil().to(torch.long)
+
+        if isinstance(ratio, int) and (k == ratio).all():
+            # If all graphs have exactly `ratio` or more than `ratio` entries,
+            # we can just pick the first entries in `perm` batch-wise:
+            index = torch.arange(batch_size, device=x.device) * max_num_nodes
+            index = index.view(-1, 1).repeat(1, ratio).view(-1)
+            index += torch.arange(ratio, device=x.device).repeat(batch_size)
+        else:
+            # Otherwise, compute indices per graph:
+            index = torch.cat([
+                torch.arange(k[i], device=x.device) + i * max_num_nodes
+                for i in range(batch_size)
+            ], dim=0)
+
+        perm = perm[index]
+        return perm
+
+    raise ValueError("At least one of the 'ratio' and 'min_score' parameters "
+                     "must be specified")
+
+
+class SelectTopK(Select):
+    # TODO (matthias) Add documentation.
+    def __init__(
+        self,
+        in_channels: int,
+        ratio: Union[int, float] = 0.5,
+        min_score: Optional[float] = None,
+        act: Union[str, Callable] = 'tanh',
+    ):
+        super().__init__()
+
+        if ratio is None and min_score is None:
+            raise ValueError(f"At least one of the 'ratio' and 'min_score' "
+                             f"parameters must be specified in "
+                             f"'{self.__class__.__name__}'")
+
+        self.in_channels = in_channels
+        self.ratio = ratio
+        self.min_score = min_score
+        self.act = activation_resolver(act)
+
+        self.weight = torch.nn.Parameter(torch.Tensor(1, in_channels))
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        uniform(self.in_channels, self.weight)
+
+    def forward(
+        self,
+        x: Tensor,
+        batch: Optional[Tensor] = None,
+    ) -> SelectOutput:
+        """"""
+        if batch is None:
+            batch = x.new_zeros(x.size(0), dtype=torch.long)
+
+        x.view(-1, 1) if x.dim() == 1 else x
+        score = (x * self.weight).sum(dim=-1)
+
+        if self.min_score is None:
+            score = self.act(score / self.weight.norm(p=2, dim=-1))
+        else:
+            score = softmax(score, batch)
+
+        node_index = topk(score, self.ratio, batch, self.min_score)
+
+        return SelectOutput(
+            node_index=node_index,
+            num_nodes=x.size(0),
+            cluster_index=torch.arange(node_index.size(0), device=x.device),
+            num_clusters=node_index.size(0),
+            weight=score[node_index],
+        )
+
+    def __repr__(self) -> str:
+        if self.min_score is None:
+            arg = f'ratio={self.ratio}'
+        else:
+            arg = f'min_score={self.min_score}'
+        return f'{self.__class__.__name__}({self.in_channels}, {arg})'
diff --git a/torch_geometric/nn/pool/topk_pool.py b/torch_geometric/nn/pool/topk_pool.py
index 998472870..661638e3a 100644
--- a/torch_geometric/nn/pool/topk_pool.py
+++ b/torch_geometric/nn/pool/topk_pool.py
@@ -2,75 +2,11 @@ from typing import Callable, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
-from torch.nn import Parameter
 
-from torch_geometric.nn.inits import uniform
-from torch_geometric.utils import scatter, softmax
+from torch_geometric.nn.pool.select import SelectTopK
 from torch_geometric.utils.num_nodes import maybe_num_nodes
 
 
-def topk(
-    x: Tensor,
-    ratio: Optional[Union[float, int]],
-    batch: Tensor,
-    min_score: Optional[float] = None,
-    tol: float = 1e-7,
-) -> Tensor:
-    if min_score is not None:
-        # Make sure that we do not drop all nodes in a graph.
-        scores_max = scatter(x, batch, reduce='max')[batch] - tol
-        scores_min = scores_max.clamp(max=min_score)
-
-        perm = (x > scores_min).nonzero().view(-1)
-
-    elif ratio is not None:
-        num_nodes = scatter(batch.new_ones(x.size(0)), batch, reduce='sum')
-        batch_size, max_num_nodes = num_nodes.size(0), int(num_nodes.max())
-
-        cum_num_nodes = torch.cat(
-            [num_nodes.new_zeros(1),
-             num_nodes.cumsum(dim=0)[:-1]], dim=0)
-
-        index = torch.arange(batch.size(0), dtype=torch.long, device=x.device)
-        index = (index - cum_num_nodes[batch]) + (batch * max_num_nodes)
-
-        dense_x = x.new_full((batch_size * max_num_nodes, ), -60000.0)
-        dense_x[index] = x
-        dense_x = dense_x.view(batch_size, max_num_nodes)
-
-        _, perm = dense_x.sort(dim=-1, descending=True)
-
-        perm = perm + cum_num_nodes.view(-1, 1)
-        perm = perm.view(-1)
-
-        if ratio >= 1:
-            k = num_nodes.new_full((num_nodes.size(0), ), int(ratio))
-            k = torch.min(k, num_nodes)
-        else:
-            k = (float(ratio) * num_nodes.to(x.dtype)).ceil().to(torch.long)
-
-        if isinstance(ratio, int) and (k == ratio).all():
-            # If all graphs have exactly `ratio` or more than `ratio` entries,
-            # we can just pick the first entries in `perm` batch-wise:
-            index = torch.arange(batch_size, device=x.device) * max_num_nodes
-            index = index.view(-1, 1).repeat(1, ratio).view(-1)
-            index += torch.arange(ratio, device=x.device).repeat(batch_size)
-        else:
-            # Otherwise, compute indices per graph:
-            index = torch.cat([
-                torch.arange(k[i], device=x.device) + i * max_num_nodes
-                for i in range(batch_size)
-            ], dim=0)
-
-        perm = perm[index]
-
-    else:
-        raise ValueError("At least one of 'min_score' and 'ratio' parameters "
-                         "must be specified")
-
-    return perm
-
-
 def filter_adj(
     edge_index: Tensor,
     edge_attr: Optional[Tensor],
@@ -157,22 +93,18 @@ class TopKPooling(torch.nn.Module):
     ):
         super().__init__()
 
-        if isinstance(nonlinearity, str):
-            nonlinearity = getattr(torch, nonlinearity)
-
         self.in_channels = in_channels
         self.ratio = ratio
         self.min_score = min_score
         self.multiplier = multiplier
-        self.nonlinearity = nonlinearity
 
-        self.weight = Parameter(torch.Tensor(1, in_channels))
+        self.select = SelectTopK(in_channels, ratio, min_score, nonlinearity)
 
         self.reset_parameters()
 
     def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""
-        uniform(self.in_channels, self.weight)
+        self.select.reset_parameters()
 
     def forward(
         self,
@@ -199,23 +131,17 @@ class TopKPooling(torch.nn.Module):
             batch = edge_index.new_zeros(x.size(0))
 
         attn = x if attn is None else attn
-        attn = attn.unsqueeze(-1) if attn.dim() == 1 else attn
-        score = (attn * self.weight).sum(dim=-1)
-
-        if self.min_score is None:
-            score = self.nonlinearity(score / self.weight.norm(p=2, dim=-1))
-        else:
-            score = softmax(score, batch)
+        select_output = self.select(attn, batch)
 
-        perm = topk(score, self.ratio, batch, self.min_score)
-        x = x[perm] * score[perm].view(-1, 1)
+        perm = select_output.node_index
+        x = x[perm] * select_output.weight.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=score.size(0))
+                                           num_nodes=select_output.num_nodes)
 
-        return x, edge_index, edge_attr, batch, perm, score[perm]
+        return x, edge_index, edge_attr, batch, perm, select_output.weight
 
     def __repr__(self) -> str:
         if self.min_score is None:


Test patch:
diff --git a/test/nn/pool/select/test_select_topk.py b/test/nn/pool/select/test_select_topk.py
new file mode 100644
index 000000000..a0418fff9
--- /dev/null
+++ b/test/nn/pool/select/test_select_topk.py
@@ -0,0 +1,57 @@
+import pytest
+import torch
+
+from torch_geometric.nn.pool.select import SelectOutput, SelectTopK
+from torch_geometric.nn.pool.select.topk import topk
+from torch_geometric.testing import is_full_test
+
+
+def test_topk_ratio():
+    x = torch.Tensor([2, 4, 5, 6, 2, 9])
+    batch = torch.tensor([0, 0, 1, 1, 1, 1])
+
+    perm1 = topk(x, 0.5, batch)
+    assert perm1.tolist() == [1, 5, 3]
+    assert x[perm1].tolist() == [4, 9, 6]
+    assert batch[perm1].tolist() == [0, 1, 1]
+
+    perm2 = topk(x, 2, batch)
+    assert perm2.tolist() == [1, 0, 5, 3]
+    assert x[perm2].tolist() == [4, 2, 9, 6]
+    assert batch[perm2].tolist() == [0, 0, 1, 1]
+
+    perm3 = topk(x, 3, batch)
+    assert perm3.tolist() == [1, 0, 5, 3, 2]
+    assert x[perm3].tolist() == [4, 2, 9, 6, 5]
+    assert batch[perm3].tolist() == [0, 0, 1, 1, 1]
+
+    if is_full_test():
+        jit = torch.jit.script(topk)
+        assert torch.equal(jit(x, 0.5, batch), perm1)
+        assert torch.equal(jit(x, 2, batch), perm2)
+        assert torch.equal(jit(x, 3, batch), perm3)
+
+
+@pytest.mark.parametrize('min_score', [None, 2.0])
+def test_select_topk(min_score):
+    if min_score is not None:
+        return
+    x = torch.randn(6, 16)
+    batch = torch.tensor([0, 0, 1, 1, 1, 1])
+
+    pool = SelectTopK(16, min_score=min_score)
+
+    if min_score is None:
+        assert str(pool) == 'SelectTopK(16, ratio=0.5)'
+    else:
+        assert str(pool) == 'SelectTopK(16, min_score=2.0)'
+
+    out = pool(x, batch)
+    assert isinstance(out, SelectOutput)
+
+    assert out.num_nodes == 6
+    assert out.num_clusters <= out.num_nodes
+    assert out.node_index.min() >= 0
+    assert out.node_index.max() < out.num_nodes
+    assert out.cluster_index.min() == 0
+    assert out.cluster_index.max() == out.num_clusters - 1
diff --git a/test/nn/pool/test_topk_pool.py b/test/nn/pool/test_topk_pool.py
index 8201db53d..a0616183f 100644
--- a/test/nn/pool/test_topk_pool.py
+++ b/test/nn/pool/test_topk_pool.py
@@ -1,35 +1,10 @@
 import torch
 
-from torch_geometric.nn.pool.topk_pool import TopKPooling, filter_adj, topk
+from torch_geometric.nn.pool import TopKPooling
+from torch_geometric.nn.pool.topk_pool import filter_adj
 from torch_geometric.testing import is_full_test
 
 
-def test_topk():
-    x = torch.Tensor([2, 4, 5, 6, 2, 9])
-    batch = torch.tensor([0, 0, 1, 1, 1, 1])
-
-    perm1 = topk(x, 0.5, batch)
-    assert perm1.tolist() == [1, 5, 3]
-    assert x[perm1].tolist() == [4, 9, 6]
-    assert batch[perm1].tolist() == [0, 1, 1]
-
-    perm2 = topk(x, 2, batch)
-    assert perm2.tolist() == [1, 0, 5, 3]
-    assert x[perm2].tolist() == [4, 2, 9, 6]
-    assert batch[perm2].tolist() == [0, 0, 1, 1]
-
-    perm3 = topk(x, 3, batch)
-    assert perm3.tolist() == [1, 0, 5, 3, 2]
-    assert x[perm3].tolist() == [4, 2, 9, 6, 5]
-    assert batch[perm3].tolist() == [0, 0, 1, 1, 1]
-
-    if is_full_test():
-        jit = torch.jit.script(topk)
-        assert torch.equal(jit(x, 0.5, batch), perm1)
-        assert torch.equal(jit(x, 2, batch), perm2)
-        assert torch.equal(jit(x, 3, batch), perm3)
-
-
 def test_filter_adj():
     edge_index = torch.tensor([[0, 0, 1, 1, 2, 2, 3, 3],
                                [1, 3, 0, 2, 1, 3, 0, 2]])