Torch Geometric Pool: the PyTorch library for pooling in Graph Neural Networks

TL;DR

Torch Geometric Pool (tgp) is a unified, flexible pooling library for Graph Neural Networks built on PyTorch Geometric, offering standardized interfaces and diverse pooling methods.

Contribution

It introduces a common software interface for graph pooling methods, enabling easier comparison and reuse across different pooling techniques.

Findings

Provides 20 hierarchical poolers with standardized outputs.

Supports dense poolers in batched and unbatched modes.

Includes workflows for caching and pre-coarsening.

Abstract

Torch Geometric Pool (tgp) is a pooling library built on top of PyTorch Geometric. Graph pooling methods differ in how they assign nodes to supernodes, how they handle batches, what they return after pooling, and whether they expose auxiliary losses. These differences make it hard to compare methods or reuse the same model code across them. tgp addresses this problem with a common software interface based on the Select-Reduce-Connect-Lift (SRCL) decomposition. The library provides 20 hierarchical poolers, standardized output objects, standalone readout modules, support for dense poolers in batched and unbatched mode, and workflows for caching and pre-coarsening. It is released under the MIT license on GitHub and PyPI, with comprehensive documentation, tutorials, and examples.