Jraph
A Graph Neural Network Library in Jax
...It provides an efficient and flexible framework for representing, manipulating, and training models on graph-structured data. The core of Jraph is the GraphsTuple data structure, which enables users to define graphs with arbitrary node, edge, and global attributes, and to batch variable-sized graphs efficiently for JAX’s just-in-time compilation. The library includes a comprehensive set of utilities for batching, padding, masking, and partitioning graph data, making it ideal for distributed and large-scale GNN experiments. Jraph also comes with a model zoo—a collection of forkable reference implementations of common message-passing GNN architectures, such as Graph Networks, Graph Convolutional Networks, and Graph Attention Networks.