jraph

What it is and what it does

Jraph is a lightweight library for graph neural networks built on top of JAX. It provides a GraphsTuple data structure—a named tuple that represents one or more directed graphs with node features, edge features, connectivity information (senders and receivers), and optional global attributes. The library includes utilities for batching multiple graphs, padding and masking for JIT compilation of variable-shaped graphs, and a model zoo with reference implementations of graph neural network architectures like GraphNetwork that use message-passing algorithms.

The package is designed for flexibility: it does not manage neural network parameters itself, leaving that to frameworks like Haiku or Flax. Instead, it provides the graph representation and message-passing infrastructure, allowing you to define custom update functions for edges, nodes, and global features. It also includes experimental support for distributing large graphs across multiple devices. However, the repository is archived and no longer maintained as of March 2024, so it may face compatibility issues with newer JAX releases.

Use it for:

• Building graph neural network models for molecular property prediction or other graph-structured data using JAX's functional programming paradigm.
• Batching and processing multiple graphs of different sizes efficiently with automatic padding and masking for JIT compilation.
• Implementing custom message-passing algorithms by defining update functions for edges, nodes, and global graph attributes.
• Distributing very large graphs (millions of edges) across multiple devices using the experimental distributed MPNN implementation.
• Prototyping graph neural network architectures with reference models from the zoo that you can fork and adapt.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Jraph provides data structures and utilities for building and working with graph neural networks in JAX, including a GraphsTuple representation for graphs and a model zoo with reference implementations of message-passing algorithms.

Yes, if you are committed to JAX for graph neural networks and can tolerate an abandoned codebase. The library is stable and well-designed, with low install friction and no known vulnerabilities. However, verify compatibility with your JAX version before adopting it for new projects, since there will be no maintenance updates. It is most suitable for research or prototyping rather than production systems requiring ongoing support.

Install

jraph on PyPI

pip install jraph

uv add jraph

poetry add jraph

Installing jraph

Before you install

Installation is straightforward with low friction. The package is abandoned as of 2024-03-18 with no recent maintenance, though it was marked Production/Stable at its last release in August 2022.

License in practice

Apache 2.0 is permissive, allowing commercial and private use with minimal restrictions beyond attribution and liability disclaimers.

Quickstart

pip install jraph

import jraph
import jax.numpy as jnp

node_features = jnp.array([[0.], [1.], [2.]])
senders = jnp.array([0, 1, 2])
receivers = jnp.array([1, 2, 0])
edges = jnp.array([[5.], [6.], [7.]])
n_node = jnp.array([3])
n_edge = jnp.array([3])

graph = jraph.GraphsTuple(nodes=node_features, senders=senders, 
                          receivers=receivers, edges=edges, 
                          n_node=n_node, n_edge=n_edge, globals=None)

Requires JAX and jaxlib; JAX has platform-specific installation steps (CPU, GPU, or TPU variants).

Verify before relying

• Whether the package remains compatible with current JAX versions given its abandonment status.
• Performance characteristics on large graphs or distributed setups beyond the experimental distributed MPNN implementation mentioned.

Package facts

Evidence: jraph-0.0.6.dev0-py3-none-any.whl

Keywords: jax, graph, neural, networks, python, machine, learning

Tags

More Python Modules packages