Jraph

Jraph (pronounced "giraffe") is a lightweight library for building Graph Neural Networks (GNNs) in JAX. It provides a fundamental data structure for graphs (`GraphsTuple`), a set of utilities for working with graphs (e.g., batching, padding, message passing), and a 'zoo' of forkable GNN models. Leveraging JAX's automatic differentiation and XLA compilation, Jraph aims for high performance and flexibility in GNN research. It is currently in a pre-release development state (v0.0.6.dev0) with frequent updates.