Skip to content

Pooling Layers

The k3_node.layers.pool module provides global readouts, hierarchical graph coarsening, cluster-based pooling, and spatial graph construction operators.

Global Graph Readout

global_add_pool

k3_node.layers.pool.global_add_pool(x, batch=None, size=None)

Returns batch-wise graph-level-outputs by adding node features across the node dimension.

Example 
import numpy as np
from k3_node.layers import global_add_pool

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1

out = global_add_pool(x, batch)  # one row per graph
print(tuple(out.shape))  # (2, 8)

global_mean_pool

k3_node.layers.pool.global_mean_pool(x, batch=None, size=None)

Returns batch-wise graph-level-outputs by averaging node features across the node dimension.

Example 
import numpy as np
from k3_node.layers import global_mean_pool

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1

out = global_mean_pool(x, batch)  # one row per graph
print(tuple(out.shape))  # (2, 8)

global_max_pool

k3_node.layers.pool.global_max_pool(x, batch=None, size=None)

Returns batch-wise graph-level-outputs by taking the channel-wise maximum across the node dimension.

Example 
import numpy as np
from k3_node.layers import global_max_pool

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1

out = global_max_pool(x, batch)  # one row per graph
print(tuple(out.shape))  # (2, 8)

Hierarchical Node Pooling

TopKPooling

k3_node.layers.pool.TopKPooling

Bases: Layer

:math:\mathrm{top}_k pooling operator from the "Graph U-Nets" <https://arxiv.org/abs/1905.05178>, "Towards Sparse Hierarchical Graph Classifiers" <https://arxiv.org/abs/1811.01287> and "Understanding Attention and Generalization in Graph Neural Networks" <https://arxiv.org/abs/1905.02850>_ papers.

Example 
import numpy as np
from k3_node.layers import TopKPooling

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1

layer = TopKPooling(in_channels=8, ratio=0.5)  # keep half of the nodes of each graph
out = layer(x, edge_index, batch=batch)
x_pool, edge_index_pool, edge_attr_pool, batch_pool = out[0], out[1], out[2], out[3]
print(tuple(x_pool.shape))  # (6, 8)

call(x, edge_index, edge_attr=None, batch=None, attn=None)

Forward pass.

reset_parameters()

Resets all learnable parameters of the module.

SAGPooling

k3_node.layers.pool.SAGPooling

Bases: Layer

The self-attention pooling operator from the "Self-Attention Graph Pooling" <https://arxiv.org/abs/1904.08082> and "Understanding Attention and Generalization in Graph Neural Networks" <https://arxiv.org/abs/1905.02850> papers.

Example 
import numpy as np
from k3_node.layers import SAGPooling

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1

layer = SAGPooling(in_channels=8, ratio=0.5)  # keep half of the nodes of each graph
out = layer(x, edge_index, batch=batch)
x_pool, edge_index_pool, edge_attr_pool, batch_pool = out[0], out[1], out[2], out[3]
print(tuple(x_pool.shape))  # (6, 8)

call(x, edge_index, edge_attr=None, batch=None, attn=None)

Forward pass.

reset_parameters()

Resets all learnable parameters of the module.

EdgePooling

k3_node.layers.pool.EdgePooling

Bases: Layer

The edge pooling operator from the "Towards Graph Pooling by Edge Contraction" <https://graphreason.github.io/papers/17.pdf> and "Edge Contraction Pooling for Graph Neural Networks" <https://arxiv.org/abs/1905.10990> papers.

Example 
import numpy as np
from k3_node.layers import EdgePooling

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1

layer = EdgePooling(in_channels=8)
x_pool, edge_index_pool, batch_pool, unpool_info = layer(x, edge_index, batch)
# The number of clusters depends on the learned edge scores
print(x_pool.shape[0] <= 10, x_pool.shape[1])  # True 8: fewer nodes, same features

call(x, edge_index, batch, training=False)

Forward pass.

unpool(x, unpool_info)

Unpools a previous edge pooling step.

ASAPooling

k3_node.layers.pool.ASAPooling

Bases: Layer

The Adaptive Structure Aware Pooling operator from the "ASAP: Adaptive Structure Aware Pooling for Learning Hierarchical Graph Representations" <https://arxiv.org/abs/1911.07979>_ paper.

Example 
import numpy as np
from k3_node.layers import ASAPooling

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1

layer = ASAPooling(in_channels=8, ratio=0.5)  # keep half of the nodes of each graph
out = layer(x, edge_index, batch=batch)
x_pool, edge_index_pool, edge_attr_pool, batch_pool = out[0], out[1], out[2], out[3]
print(tuple(x_pool.shape))  # (6, 8)

call(x, edge_index, edge_weight=None, batch=None, training=False)

Forward pass.

PANPooling

k3_node.layers.pool.PANPooling

Bases: Layer

The path integral based pooling operator from the "Path Integral Based Convolution and Pooling for Graph Neural Networks" <https://arxiv.org/abs/2006.16811>_ paper.

Example 
import numpy as np
from k3_node.layers import PANPooling, PANConv

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1

# PANPooling consumes the path weights produced by PANConv
x, weights = PANConv(in_channels=8, out_channels=8, filter_size=2)(x, edge_index)
layer = PANPooling(in_channels=8, ratio=0.5)
x_pool, edge_index_pool, edge_attr_pool, batch_pool, perm, score = layer(x, weights, batch=batch)
print(tuple(x_pool.shape))  # (6, 8)

call(x, M, batch=None)

Forward pass.

Parameters:

Name Type Description Default
x

Node feature matrix.

 required
M

MET matrix, either a tuple/list (edge_index, edge_weight) or an object with .coo() method (like PyG SparseTensor).

 required
batch Optional[any]

Batch vector.

 None

MemPooling

k3_node.layers.pool.MemPooling

Bases: Layer

Memory based pooling layer from "Memory-Based Graph Networks" <https://arxiv.org/abs/2002.09518>_ paper.

Example 
import numpy as np
from k3_node.layers import MemPooling

x = np.random.rand(2, 5, 8).astype("float32")  # dense batch: 2 graphs with 5 nodes each
layer = MemPooling(in_channels=8, out_channels=16, heads=2, num_clusters=3)
x_pool, assignment = layer(x)
print(tuple(x_pool.shape))  # (2, 3, 16)
print(tuple(assignment.shape))  # (2, 5, 3): soft assignment of nodes to clusters

call(x, batch=None, mask=None, max_num_nodes=None, batch_size=None)

Forward pass.

kl_loss(S) staticmethod

The additional KL divergence-based loss.

ClusterPooling

k3_node.layers.pool.ClusterPooling

Bases: Layer

The cluster pooling operator from the "Edge-Based Graph Component Pooling" <https://arxiv.org/abs/2409.11856>_ paper.

Example 
import numpy as np
from k3_node.layers import ClusterPooling

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1

layer = ClusterPooling(in_channels=8)
x_pool, edge_index_pool, batch_pool, unpool_info = layer(x, edge_index, batch)
# The number of clusters depends on the learned edge scores
print(x_pool.shape[0] <= 10, x_pool.shape[1])  # True 8: fewer nodes, same features

call(x, edge_index, batch, training=False)

Forward pass.

Spatial & Neighborhood Pooling

avg_pool_neighbor_x

k3_node.layers.pool.avg_pool_neighbor_x(data, edge_index=None, flow='source_to_target')

Average-pools neighboring node features.

Example 
import numpy as np
from k3_node.layers import avg_pool_neighbor_x

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

out = avg_pool_neighbor_x(x, edge_index=edge_index)  # pool each node with its neighbors
print(tuple(out.shape))  # (10, 8)

max_pool_neighbor_x

k3_node.layers.pool.max_pool_neighbor_x(data, edge_index=None, flow='source_to_target')

Max-pools neighboring node features.

Example 
import numpy as np
from k3_node.layers import max_pool_neighbor_x

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

out = max_pool_neighbor_x(x, edge_index=edge_index)  # pool each node with its neighbors
print(tuple(out.shape))  # (10, 8)

voxel_grid

k3_node.layers.pool.voxel_grid

voxel_grid(pos, size, batch=None, start=None, end=None)

Voxel grid pooling that clusters points within the same voxel.

Example 
import numpy as np
from k3_node.layers import voxel_grid

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
pos = np.random.rand(10, 3).astype("float32")  # 3D positions

cluster = voxel_grid(pos, size=0.5)  # voxel id of each point
print(tuple(cluster.shape))  # (10,)

fps

k3_node.layers.pool.fps(x, batch=None, ratio=0.5, random_start=True, batch_size=None)

Farthest Point Sampling algorithm.

Example 
import numpy as np
from k3_node.layers import fps

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
pos = np.random.rand(10, 3).astype("float32")  # 3D positions
batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1

index = fps(pos, batch, ratio=0.4)  # farthest point sampling: 40% of each graph's points
print(tuple(index.shape))  # (4,)

graclus

k3_node.layers.pool.graclus

graclus(edge_index, weight=None, num_nodes=None)

A greedy clustering algorithm of picking an unmarked vertex and matching it with one of its unmarked neighbors that maximizes its edge weight.

Example 
import numpy as np
from k3_node.layers import graclus

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges

cluster = graclus(edge_index, num_nodes=10)  # greedy matching of neighboring nodes
print(tuple(cluster.shape))  # (10,): a cluster id for every node

Graph Construction

radius_graph

k3_node.layers.pool.radius_graph(x, r, batch=None, loop=False, max_num_neighbors=32, flow='source_to_target', num_workers=1, batch_size=None)

Computes graph edges to all points within a given distance r.

Example 
import numpy as np
from k3_node.layers import radius_graph

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
pos = np.random.rand(10, 3).astype("float32")  # 3D positions

edge_index = radius_graph(pos, r=0.5)  # connect points closer than 0.5
print(edge_index.shape[0])  # 2

knn_graph

k3_node.layers.pool.knn_graph(x, k, batch=None, loop=False, flow='source_to_target', cosine=False, num_workers=1, batch_size=None)

Computes graph edges to the nearest k points.

Example 
import numpy as np
from k3_node.layers import knn_graph

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
pos = np.random.rand(10, 3).astype("float32")  # 3D positions

edge_index = knn_graph(pos, k=3)  # connect every point to its 3 nearest neighbors
print(tuple(edge_index.shape))  # (2, 30)

Unpooling

Interpolates features back from a coarsened point set/graph onto a denser one, mirroring torch_geometric.nn.unpool.

knn_interpolate

k3_node.layers.unpool.knn_interpolate

knn_interpolate(x, pos_x, pos_y, batch_x=None, batch_y=None, k=3, num_workers=1)

The k-NN interpolation from the "PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space" <https://arxiv.org/abs/1706.02413>_ paper.

For each point :math:y with position :math:\mathbf{p}(y), its interpolated features :math:\mathbf{f}(y) are given by

.. math:: \mathbf{f}(y) = \frac{\sum_{i=1}^k w(x_i) \mathbf{f}(x_i)}{\sum_{i=1}^k w(x_i)} \textrm{, where } w(x_i) = \frac{1}{d(\mathbf{p}(y), \mathbf{p}(x_i))^2}

and :math:\{ x_1, \ldots, x_k \} denoting the :math:k nearest points to :math:y.

Parameters:

Name Type Description Default
x

Node feature matrix :math:\mathbf{X} \in \mathbb{R}^{N \times F}.

 required
pos_x

Node position matrix :math:\in \mathbb{R}^{N \times d}.

 required
pos_y

Upsampled node position matrix :math:\in \mathbb{R}^{M \times d}.

 required
batch_x

Batch vector assigning each node from :math:\mathbf{X} to a specific example. (default: :obj:None)

 None
batch_y

Batch vector assigning each node from :math:\mathbf{Y} to a specific example. (default: :obj:None)

 None
k int

Number of neighbors. (default: :obj:3)

 3
num_workers int

Unused, kept for API compatibility.

 1
Example 
import numpy as np
from k3_node.layers import knn_interpolate

x = np.random.rand(6, 8).astype("float32")  # features of 6 coarse points
pos_x = np.random.rand(6, 3).astype("float32")
pos_y = np.random.rand(10, 3).astype("float32")  # 10 fine points to interpolate onto
out = knn_interpolate(x, pos_x, pos_y, k=3)
print(tuple(out.shape))  # (10, 8)

Regularization Functionals

Differentiable graph-pooling regularizers from k3_node.layers.functional, typically added as auxiliary loss terms alongside a pooling layer (e.g. DMoNPooling, MinCutPooling) to encourage balanced, well-separated cluster assignments.

bro

k3_node.layers.functional.bro

bro(x, batch, p=2)

The Batch Representation Orthogonality penalty from the "Improving Molecular Graph Neural Network Explainability with Orthonormalization and Induced Sparsity" <https://arxiv.org/abs/2105.04854>_ paper.

Computes a regularization for each graph representation in a mini-batch according to

.. math:: \mathcal{L}_{\textrm{BRO}}^\mathrm{graph} = || \mathbf{HH}^T - \mathbf{I}||_p

and returns an average over all graphs in the batch.

Example 
import numpy as np
from k3_node.layers import bro

x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
batch = np.repeat([0, 1], 5)  # nodes 0-4 belong to graph 0, nodes 5-9 to graph 1

loss = bro(x, batch)  # Batch Representation Orthogonality penalty
print(tuple(loss.shape))  # (): a scalar

gini

k3_node.layers.functional.gini

gini(w)

The Gini coefficient from the "Improving Molecular Graph Neural Network Explainability with Orthonormalization and Induced Sparsity" <https://arxiv.org/abs/2105.04854>_ paper.

Computes a regularization penalty :math:\in [0, 1] for each row of a matrix according to

.. math:: \mathcal{L}\textrm{Gini}^i = \sum_j^n \sum - w_{ij'}|}{2 (n^2 - n)\bar{w_i}}}^n \frac{|w_{ij

and returns an average over all rows.

Parameters:

Name Type Description Default
w

A two-dimensional tensor.

 required
Example 
import numpy as np
from k3_node.layers import gini

w = np.random.rand(16, 8).astype("float32")  # e.g. a weight matrix
print(tuple(gini(w).shape))  # (): scalar sparsity measure