Install & Compatibility
Where this runs
tested against v0.0.6.dev0 · pip install
no network on importno background threads
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
py 3.10
✕ build_error
4/8 runs
py 3.11
✕ build_error
4/8 runs
py 3.12
✕ build_error
4/8 runs
py 3.13
✕ build_error
4/8 runs
py 3.9
✕ build_error
4/8 runs
Code
Verified usage
Verified import paths — ran on the pinned version, not inferred.
GraphsTuple
✓ from jraph import GraphsTuple
GraphNetwork
✓ from jraph.models import GraphNetwork
Models are typically imported from the `jraph.models` submodule, or often constructed directly from custom update functions.
batch
✓ from jraph import batch
Utilities like `batch`, `unbatch`, `pad_with_graphs` are directly available under the `jraph` namespace.
This quickstart demonstrates how to construct a basic `GraphsTuple` object, which is the core data structure for representing graphs in Jraph. It defines nodes, edges, and optional global features using `jax.numpy` arrays, along with the necessary `senders`, `receivers`, `n_node`, and `n_edge` arrays to describe the graph structure.
import jraph
import jax.numpy as jnp
# Define node features, 3 nodes, each with a scalar feature
node_features = jnp.array([[0.], [1.], [2.]])
# Define edges: 0 -> 1, 1 -> 2
senders = jnp.array([0, 1])
receivers = jnp.array([1, 2])
# Edge features (optional), 2 edges, each with a scalar feature
edge_features = jnp.array([[10.], [20.]])
# Global features (optional), 1 graph, with a scalar feature
global_features = jnp.array([[100.]])
# Number of nodes and edges per graph (for a single graph)
n_node = jnp.array([len(node_features)])
n_edge = jnp.array([len(senders)])
# Create a GraphsTuple
graph = jraph.GraphsTuple(
nodes=node_features,
edges=edge_features,
receivers=receivers,
senders=senders,
globals=global_features,
n_node=n_node,
n_edge=n_edge
)
print(graph)
print(f"Nodes: {graph.nodes.shape}, Edges: {graph.edges.shape}")
Debug
Known issues
breakingJraph is currently in early development (0.0.x.dev0 versions), meaning API changes can occur frequently and without adherence to semantic versioning for breaking changes. Code developed with one minor version may not be compatible with the next.fixPin Jraph to a specific version (`pip install jraph==0.0.6.dev0`) and regularly review the GitHub changelog for updates when upgrading.
affects: All 0.0.x.dev0 versions
gotchaJraph focuses on graph data structures and message passing, but it does not manage parameters for graph neural networks. Users need to integrate with external JAX-native neural network libraries like Haiku or Flax for parameter management and model construction.fixFamiliarize yourself with Haiku or Flax for defining and managing model parameters within your Jraph-based GNNs. See examples for integration patterns.
affects: All versions
gotchaThe `jraph.unbatch` utility does not support `jax.jit` compilation because the shapes of the unbatched output graphs are data-dependent, preventing JAX from tracing a static computation graph.fixAvoid using `jraph.unbatch` inside `jax.jit` decorated functions. Unbatch graphs outside of jitted functions, or consider using padding and masking techniques for variable-sized graphs within jitted contexts (e.g., `pad_with_graphs`).
affects: All versions
gotchaIndexing for nodes and edges within a batched `GraphsTuple` is absolute (cumulative) across all graphs in the batch, rather than relative to each individual graph. This can lead to off-by-one errors or incorrect feature access if not handled carefully.fixAlways account for the `n_node` and `n_edge` properties of the `GraphsTuple` to calculate the correct absolute indices when accessing or manipulating features for a specific graph within a batch. For example, the nodes of the i-th graph start at `sum(graph.n_node[:i])`.
affects: All versions
Errors
Common errors & fixes
AttributeError: module 'jraph' has no attribute 'GraphsTuple'
This usually indicates an outdated `jraph` installation or a confusion with older API patterns. In some environments, `jraph` might be shadowed or incorrectly installed.
fixEnsure `jraph` is correctly installed and up-to-date (`pip install --upgrade jraph`). Verify the import statement is `from jraph import GraphsTuple`.
ValueError: Found graph bigger than batch size. Valid Batch Size: {...}, Graph Size: {...}
This error typically occurs during batching or padding operations when trying to create a batch of graphs where a single graph exceeds the maximum allowed size configured for the batch (e.g., during `pad_with_graphs`).
fixReview the parameters used for `jraph.pad_with_graphs` or similar batching utilities. Adjust the padding limits or ensure that the input graphs do not individually exceed the capacity of the padded batch. This is often related to `n_node_total` or `n_edge_total` arguments.
TypeError: 'dict' object is not callable (or similar error when passing dicts as features)
While `GraphsTuple` supports `ArrayTrees` (including nested dictionaries) for `nodes`, `edges`, and `globals`, some internal Jraph functions or external JAX operations might expect flat arrays or specific structures. Passing complex dictionary structures where a flat array is expected can cause errors.
fixEnsure that if you are using dictionary features, the functions processing these features are designed to handle `ArrayTrees` (e.g., using `jax.tree_util` functions). For simpler cases, convert dictionary features to concatenated arrays if only a single feature vector is needed per node/edge/graph for a given operation.
Upgrade
Version history
0.0.6.dev0latest on PyPI · released Aug 12, 2022
Audit
Dependencies
jaxrequiredJraph is built on JAX for high-performance numerical computation and automatic differentiation.
jaxlibrequiredJAX's compiled operations require jaxlib for the backend.
frozendictrequiredRequired for graph data structures, added as a dependency in v0.0.3.dev0.