Connect¶

`Connect`	An abstract base class implementing the \(\texttt{connect}\) operator.
`SparseConnect`	The \(\texttt{connect}\) operator for sparse methods where each node is assigned at most one supernode.
`DenseConnect`	The \(\texttt{connect}\) operator for dense pooling methods.
`EigenPoolConnect`	The \(\texttt{connect}\) operator for EigenPooling.
`KronConnect`	The \(\texttt{connect}\) operator based on Kron reduction proposed in the paper "Hierarchical Representation Learning in Graph Neural Networks with Node Decimation Pooling" (Bianchi et al., TNNLS 2020).
`sparse_connect`	Connects the nodes in the coarsened graph.

class Connect(*args, **kwargs)[source]¶

An abstract base class implementing the \(\texttt{connect}\) operator.

Specifically, \(\texttt{connect}\) determines for each pair of supernodes the presence or absence of an edge based on the existing edges between the nodes in the two supernodes.

forward(edge_index: Tensor | SparseTensor, so: SelectOutput, *, edge_weight: Tensor | None = None, **kwargs) → Tuple[Tensor | SparseTensor, Tensor | None][source]¶

Forward pass.

Parameters:

edge_index (torch.Tensor) – The original edge indices.
so (SelectOutput) – The output of the \(\texttt{select}\) operator.
edge_weight (torch.Tensor, optional) – The original edge weights. (default: None)

class SparseConnect(reduce_op: Literal['sum', 'mean', 'min', 'max', 'mul'] = 'sum', remove_self_loops: bool = True, edge_weight_norm: bool = False, degree_norm: bool = False)[source]¶

The \(\texttt{connect}\) operator for sparse methods where each node is assigned at most one supernode. This is, for example, the case of one-over-\(K\) methods such as GraclusSelect, NDPSelect, and KMISSelect.

It also works for scoring-based methods such as TopkSelect that compute the pooled adjacency as

\[\mathbf{A}_{\text{pool}} = \mathbf{A}_{\mathbf{i},\mathbf{i}},\]

where \(\mathbf{i}\) denotes the set of supernodes.

Parameters:

reduce_op (ConnectionType, optional) – The aggregation function to be applied to nodes in the same cluster. Can be any string admitted by scatter (e.g., 'sum', 'mean', 'max') or any ConnectionType. (default: sum)
remove_self_loops (bool, optional) – Whether to remove self-loops from the graph after coarsening. (default: True)
edge_weight_norm (bool, optional) – Whether to normalize the edge weights by dividing by the maximum absolute value per graph. (default: False)
degree_norm (bool, optional) – If True, the adjacency matrix will be symmetrically normalized using \(D^{-1/2} A D^{-1/2}\) where \(D\) is the degree matrix. (default: False)

Forward pass.

Parameters:

edge_index (Adj) – The connectivity matrix. It can either be a torch_sparse.SparseTensor of (sparse) shape \([N, N]\), where \(N\) is the number of nodes in the batch or a Tensor of shape \([2, E]\), where \(E\) is the number of edges in the batch.
so (SelectOutput) – The output of the \(\texttt{select}\) operator.
edge_weight (Tensor, optional) – A vector of shape \([E]\) containing the weights of the edges. (default: None)
batch_pooled (Tensor, optional) – Batch vector which assigns each supernode to a specific graph. Required when edge_weight_norm=True for per-graph normalization. (default: None)

Returns:

The pooled adjacency matrix and the edge weights. If the pooled adjacency is a torch_sparse.SparseTensor, returns None as the edge weights.

Return type:

(Adj, Tensor or None)

class DenseConnect(remove_self_loops: bool = True, degree_norm: bool = True, adj_transpose: bool = True, edge_weight_norm: bool = False, sparse_output: bool = False)[source]¶

The \(\texttt{connect}\) operator for dense pooling methods.

Input representations:

Batched dense inputs: adjacency \([B, N, N]\), assignment \([B, N, K]\).
Unbatched sparse inputs: sparse adjacency and dense assignment \([N, K]\) (or \([1, N, K]\)).

Output representations:

Batched dense inputs always return a dense adjacency \([B, K, K]\) (edge weights are None).
Unbatched sparse inputs return either a dense adjacency \([B, K, K]\) or a block-diagonal sparse adjacency \([B*K, B*K]\) depending on sparse_output.

It computes the pooled adjacency matrix as:

\[\mathbf{A}_{\mathrm{pool}} = \mathbf{S}^{\top}\mathbf{A}\mathbf{S}\]

Parameters:

remove_self_loops (bool, optional) – Whether to remove self-loops from the graph after coarsening. (default: True)
degree_norm (bool, optional) – If True, the adjacency matrix will be symmetrically normalized. (default: True)
adj_transpose (bool, optional) – If True, it returns a transposed pooled adjacency matrix for batched dense outputs, so that it can be passed “as is” to the dense message passing layers. This only applies to batched dense inputs. (default: True)
edge_weight_norm (bool, optional) – Whether to normalize the edge weights by dividing by the maximum absolute value per graph. (default: False)
sparse_output (bool, optional) – Controls the output format only for unbatched inputs. If True, return a block-diagonal sparse adjacency of shape \([B*K, B*K]\). If False, return a dense adjacency of shape \([B, K, K]\). Batched dense inputs always return a dense adjacency. (default: False)

dense_connect(adj: Tensor, s: Tensor) → Tensor[source]¶

Public method to compute \(\mathbf{S}^{\top}\mathbf{A}\mathbf{S}\).

This method is a convenience wrapper for poolers that need to compute the raw pooled adjacency before applying post-processing.

Parameters:

adj (Tensor) – The dense adjacency matrix of shape \((B, N, N)\).
s (Tensor) – The dense assignment matrix of shape \((B, N, K)\).

Returns:

The pooled adjacency matrix of shape \((B, K, K)\).

Return type:

Tensor

Forward pass.

Parameters:

edge_index (Adj) – For batched dense inputs, a tensor of shape \([B, N, N]\). For unbatched sparse inputs, a sparse connectivity matrix in one of the formats supported by Adj.
so (SelectOutput) – The output of the \(\texttt{select}\) operator. The assignment matrix so.s must be a dense tensor.
edge_weight (Tensor, optional) – A vector of shape \([E]\) or \([E, 1]\) containing the weights of the edges for unbatched inputs. (default: None)
batch (Tensor, optional) – The batch vector for unbatched inputs. (default: None)
batch_pooled (Tensor, optional) – The pooled batch vector, required for edge weight normalization with unbatched sparse outputs. (default: None)

Returns:

The pooled adjacency matrix and the edge weights. If the pooled adjacency is dense, returns None for the edge weights.

Return type:

(Adj, Tensor or None)

class EigenPoolConnect(remove_self_loops: bool = True, degree_norm: bool = True, adj_transpose: bool = True, edge_weight_norm: bool = False, sparse_output: bool = False)[source]¶

The \(\texttt{connect}\) operator for EigenPooling.

EigenPooling replaces the standard \(\mathbf{S}^{\top}\mathbf{A}\mathbf{S}\) connection with:

\[\mathbf{A}_{\text{coar}} = \boldsymbol{\Omega}^{\top}\mathbf{A}_{\text{ext}}\boldsymbol{\Omega},\]

where:

\[\begin{split}\mathbf{A}_{\text{ext}} = \mathbf{A} - \mathbf{A}_{\text{int}}, \qquad (\mathbf{A}_{\text{int}})_{ij} = \begin{cases} \mathbf{A}_{ij} & \text{if } c_i = c_j \\ 0 & \text{otherwise} \end{cases}\end{split}\]

and \(\boldsymbol{\Omega}\) is the hard cluster membership matrix returned by EigenPoolSelect (i.e., so.s), with \(c_i = \arg\max_k \Omega_{ik}\).

Input representations:

Batched dense inputs: adjacency \([B, N, N]\), assignment \([B, N, K]\).
Unbatched sparse inputs: sparse adjacency and dense assignment \([N, K]\) (or \([1, N, K]\)).

Output representations:

Batched dense inputs always return a dense adjacency \([B, K, K]\) (edge weights are None).
Unbatched sparse inputs return either a dense adjacency \([B, K, K]\) or a block-diagonal sparse adjacency \([B*K, B*K]\) depending on sparse_output.

Parameters:

remove_self_loops (bool, optional) – Whether to remove self-loops after coarsening. (default: True)
degree_norm (bool, optional) – If True, symmetrically normalize the pooled adjacency. (default: True)
adj_transpose (bool, optional) – If True, transpose the dense pooled adjacency for message passing. Only applies to batched dense inputs. (default: True)
edge_weight_norm (bool, optional) – Whether to normalize edge weights by their max absolute value per graph. (default: False)
sparse_output (bool, optional) – Controls the output format only for unbatched inputs. If True, return a block-diagonal sparse adjacency of shape \([B*K, B*K]\). If False, return a dense adjacency of shape \([B, K, K]\). (default: False)

Forward pass.

Parameters:

edge_index (Adj) – For batched dense inputs, a dense adjacency tensor of shape \([B, N, N]\). For unbatched inputs, sparse connectivity in any format supported by Adj.
so (SelectOutput) – Output of the \(\texttt{select}\) operator. The assignment matrix so.s must be dense.
edge_weight (Tensor, optional) – Edge weights associated with edge_index for sparse inputs. (default: None)
batch (Tensor, optional) – Batch vector for sparse multi-graph inputs. If None, all nodes are treated as belonging to a single graph. (default: None)
batch_pooled (Tensor, optional) – Batch vector for pooled nodes, used by sparse post-processing when per-graph edge-weight normalization is enabled. (default: None)

Returns:

The coarsened adjacency and edge weights. Dense outputs return None for edge weights.

Return type:

(Adj, Tensor or None)

The coarsened adjacency is computed as:

\[\mathbf{A}_{\text{coar}} = \boldsymbol{\Omega}^{\top}\mathbf{A}_{\text{ext}}\boldsymbol{\Omega},\]

where \(\boldsymbol{\Omega} = \texttt{so.s}\).

class KronConnect(sparse_threshold: float = 0.01)[source]¶

The \(\texttt{connect}\) operator based on Kron reduction proposed in the paper “Hierarchical Representation Learning in Graph Neural Networks with Node Decimation Pooling” (Bianchi et al., TNNLS 2020).

Given two disjoint sets of nodes, \(\mathcal{V}^+\) and \(\mathcal{V}^-\), a pair of nodes \(i\) and \(j\) are connected if \(i,j \in \mathcal{V}^+\), and there is a path in the original graph that connects \(i\) and \(j\), for each other node \(k\) on the path \(k \notin \mathcal{V}^-\)

Parameters:: sparse_threshold (float, optional) – Deletes edges whose weight is inferior to the given value. (default: 1e-2)

forward(edge_index: Tensor | SparseTensor, so: SelectOutput, edge_weight: Tensor | None = None, **kwargs) → Tuple[Tensor | SparseTensor, Tensor | None][source]¶

The forward pass.

Parameters:

edge_index (Tensor) – A tensor of shape \([2, E]\), where \(E\) is the number of edges in the batch.
so (SelectOutput, optional) – The output of the \(\texttt{select}\) operator. (default: None)
edge_weight (Tensor, optional) – A vector of shape \([E]\) containing the weights of the edges. (default: None)

Returns:

The pooled adjacency matrix and the edge weights. If the pooled adjacency is a torch_sparse.SparseTensor, returns None as the edge weights.

Return type:

(Adj, Tensor or None)

sparse_connect(edge_index: Tensor | SparseTensor, edge_weight: Tensor | None = None, node_index: Tensor | None = None, cluster_index: Tensor | None = None, num_nodes: int | None = None, num_supernodes: int | None = None, remove_self_loops: bool = True, reduce_op: Literal['sum', 'mean', 'min', 'max', 'mul'] = 'sum', edge_weight_norm: bool = False, batch_pooled: Tensor | None = None, degree_norm: bool = False) → Tuple[Tensor | SparseTensor, Tensor | None][source]¶: Connects the nodes in the coarsened graph.