Download
Jraph (pronounced “giraffe”) is a lightweight JAX library developed by Google DeepMind for building and experimenting with graph neural networks (GNNs). 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.
Features
- Lightweight GraphsTuple data structure for flexible graph representation
- Distributed message-passing support for massive graphs across multiple devices
- Utilities for batching, masking, and padding to handle variable-sized graphs
- Modular model zoo of reusable graph neural network architectures
- Educational Colab tutorials and large-scale dataset examples (e.g., OGBG-MOLPCBA)
- Fully JAX-compatible for jit compilation, pmap parallelization, and scalability
Project Samples
