Jraph

Common errors

• AttributeError: module 'jraph' has no attribute 'GraphsTuple'

  cause This usually indicates an outdated `jraph` installation or a confusion with older API patterns. In some environments, `jraph` might be shadowed or incorrectly installed.
  fix Ensure `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: {...}

  cause 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`).
  fix Review 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)

  cause 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.
  fix Ensure 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.

Warnings

• breaking Jraph 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.
  Fix: Pin Jraph to a specific version (`pip install jraph==0.0.6.dev0`) and regularly review the GitHub changelog for updates when upgrading.
• gotcha Jraph 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.
  Fix: Familiarize yourself with Haiku or Flax for defining and managing model parameters within your Jraph-based GNNs. See examples for integration patterns.
• gotcha The `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.
  Fix: Avoid 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`).
• gotcha Indexing 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.
  Fix: Always 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])`.

Install

• pip install jraph Install Jraph from PyPI
• pip install git+git://github.com/deepmind/jraph.git Install latest from GitHub

Imports

• 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.

Quickstart

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}")