jraph
Jraph: A library for Graph Neural Networks in Jax
Decision gist · record as of 2026-08-14
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.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires JAX and jaxlib; JAX has platform-specific installation steps (CPU, GPU, or TPU variants).
- 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 · maintenance · safety
Apache 2.0 (permissive) — Apache 2.0 is permissive, allowing commercial and private use with minimal restrictions beyond attribution and liability disclaimers.
last release 2022-08-12 (1463 days) · last repo commit 2024-03-18 · 1,473 stars · archived
0 known vulnerabilities (OSV.dev, 2026-08-14) · 88,585 downloads/mo, #13,713 on PyPI
Alternatives
Verify before relying
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)- 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.
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 on it.
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
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.
Requires JAX and jaxlib; JAX has platform-specific installation steps (CPU, GPU, or TPU variants).
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)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
| License | Apache 2.0 permissive |
| Python support | Supports the current Python release >=3.6 |
| Install friction | Low. Pure-Python wheel |
| Runtime dependencies | 3 packagesjaxjaxlibnumpy |
| Maintenance | Abandoned 1,463 days since the last release |
| Last repo commit | repository archived |
| First released | |
| Downloads | 88,585 / month, #13,713 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
| Classifiers | Development Status :: 5 - Production/StableEnvironment :: ConsoleIntended Audience :: DevelopersIntended Audience :: Science/ResearchLicense :: OSI Approved :: Apache Software LicenseOperating System :: OS IndependentProgramming Language :: PythonProgramming Language :: Python :: 3Topic :: Scientific/Engineering :: Artificial IntelligenceTopic :: Software Development :: Libraries :: Python Modules |
Evidence: jraph-0.0.6.dev0-py3-none-any.whl
Tags
Let your AI agent find packages like this
Example. Real query, live index.
You found this page by searching. An agent finds it by wishing: SkillFed indexes 14,416 PyPI packages by what they can do, searchable in plain language.
wish › “graph neural networks jax”
- jraphJraph provides data structures and utilities for building and working…
- equinoxEquinox provides neural network and model building on top of JAX with…
- e3nn-jaxImplements equivariant convolutional neural networks for 3D…
Give your agent the search over MCP, or paste the wish link into any chat.
More Python Modules packages
Converts domain names between Unicode and ASCII-compatible encoding (Punycode) according to IDNA 2008 and Unicode Technical Standard 46, with security validation and broader script coverage than the standard library.
Install it if you work with internationalized domain names, need to validate domains, or use HTTP clients that depend on it transitively.
Setuptools is a Python build backend and package management tool that handles building, distributing, and installing Python packages, including support for C/C++ extension modules.
PyYAML parses and emits YAML 1.1 data format, enabling serialization and deserialization of configuration files and Python objects to and from human-readable YAML text.
Pydantic validates Python data structures against type hints, coercing and checking input at runtime to ensure it matches a declared schema.
Provides reusable metadata objects for use with PEP-593 `typing.Annotated` to express common constraints like bounds, collection sizes, and predicates on types.
Install it if you use or build libraries that need to express type constraints in a standardized, inspectable way—or if you want to annotate your own types with…
Provides runtime tools to inspect and introspect Python type annotations, enabling programmatic examination of type hints at execution time.
See also rax · dm-haiku · Penman · jmp · torch-geometric · flax · equinox · distrax · optax · e3nn-jax