layer_top_k_pool: TopKPool
In rdinnager/rspektral: What the Package Does (One Line, Title Case)

\loadmathjax

A gPool/Top-K layer as presented by Gao & Ji (2019) and Cangea et al. (2018).

Mode: single, disjoint.

This layer computes the following operations:

\mjdeqn\boldsymbol

y = \frac\boldsymbolX\p\|\boldsymbolp\|; \;\;\;\;\boldsymboli = \textrmrank(\boldsymboly, K); \;\;\;\;\boldsymbolX' = (\boldsymbolX \odot \textrmtanh(\boldsymboly))_\boldsymboli; \;\;\;\;\boldsymbolA' = \boldsymbolA _ \boldsymboli, \boldsymboli

where \mjeqn \textrmrank(\boldsymboly, K) returns the indices of the top K values of \mjeqn\boldsymboly, and \mjeqn\boldsymbolp is a learnable parameter vector of size \mjeqnF. \mjeqnK is defined for each graph as a fraction of the number of nodes. Note that the the gating operation \mjeqn\textrmtanh(\boldsymboly) (Cangea et al.) can be replaced with a sigmoid (Gao & Ji).

This layer temporarily makes the adjacency matrix dense in order to compute \mjeqn\boldsymbolA'. If memory is not an issue, considerable speedups can be achieved by using dense graphs directly. Converting a graph from sparse to dense and back to sparse is an expensive operation.

Input

Node features of shape (N, F);
Binary adjacency matrix of shape (N, N);
Graph IDs of shape (N, ) (only in disjoint mode);

Output

Reduced node features of shape (ratio * N, F);
Reduced adjacency matrix of shape (ratio * N, ratio * N);
Reduced graph IDs of shape (ratio * N, ) (only in disjoint mode);
If return_mask=True, the binary pooling mask of shape (ratio * N, ).

layer_top_k_pool(
  object,
  ratio,
  return_mask = FALSE,
  sigmoid_gating = FALSE,
  kernel_initializer = "glorot_uniform",
  kernel_regularizer = NULL,
  kernel_constraint = NULL,
  ...
)

`ratio`	float between 0 and 1, ratio of nodes to keep in each graph
`return_mask`	boolean, whether to return the binary mask used for pooling
`sigmoid_gating`	boolean, use a sigmoid gating activation instead of a tanh
`kernel_initializer`	initializer for the weights
`kernel_regularizer`	regularization applied to the weights
`kernel_constraint`	constraint applied to the weights