update

torch_geometric/nn/pool/connect/__init__.py

@@ -1,5 +1,6 @@
-from .base import Connect
+from .base import Connect, ConnectOutput
 
 __all__ = [
     'Connect',
+    'ConnectOutput',
 ]

torch_geometric/nn/pool/connect/base.py

@@ -50,9 +50,9 @@ def __init__(
 
 
 class Connect(torch.nn.Module):
-    r"""An abstract base class implementing custom edge connection operators as
-    described in the `"Understanding Pooling in Graph Neural Networks"
-    <https://arxiv.org/abs/1905.05178>`_ paper.
+    r"""An abstract base class for implementing custom edge connection
+    operators as described in the `"Understanding Pooling in Graph Neural
+    Networks" <https://arxiv.org/abs/1905.05178>`_ paper.
 
     Specifically, :class:`Connect` determines for each pair of supernodes the
     presence or abscene of an edge based on the existing edges between the
@@ -73,8 +73,7 @@ def forward(
     ) -> Tuple[Tensor, Optional[Tensor]]:
         r"""
         Args:
-            cluster (SelectOutput): The output of `Select`, with a mapping from
-                nodes to clusters.
+            select_output (SelectOutput): The output of :class:`Select`.
             edge_index (torch.Tensor): The edge indices.
             edge_attr (torch.Tensor, optional): The edge features.
                 (default: :obj:`None`)