jraphx.nn.models

This module contains pre-built GNN model architectures that can be used out-of-the-box for common graph learning tasks.

Pre-built GNN Models

These models provide complete architectures with multiple layers, normalization, dropout, and optional features like JumpingKnowledge connections.

GCN

class GCN(*args: Any, **kwargs: Any)[source]

Bases: BasicGNN

Graph Convolutional Network.

From “Semi-supervised Classification with Graph Convolutional Networks” https://arxiv.org/abs/1609.02907

Uses GCNConv layers for message passing.

Parameters:

  • in_features (int) – Size of input features

  • hidden_features (int) – Size of hidden layers

  • num_layers (int) – Number of GCN layers

  • out_features (int | None, default: None) – Size of output (if None, uses hidden_features)

  • dropout_rate (float, default: 0.0) – Dropout probability

  • act (Callable | None, default: <jax._src.custom_derivatives.custom_jvp object at 0x7f10e9a36ba0>) – Non-linear activation function, or None to disable the activation entirely (default: jax.nn.relu)

  • act_first (bool, default: False) – If True, apply activation before normalization

  • norm (str | None, default: None) – Normalization type (‘batch_norm’, ‘layer_norm’, ‘graph_norm’, None)

  • jk (str | None, default: None) – Jumping Knowledge mode (‘last’, ‘cat’, ‘max’, ‘lstm’, None)

  • residual (bool, default: False) – Whether to use residual connections

  • improved – Use improved GCN normalization

  • cached – Reuse a precomputed normalization for a static graph. The cache of every layer must be filled eagerly with precompute_norm() before the first forward pass; a forward pass with an empty cache raises RuntimeError.

  • add_self_loops – Add self-loops to the graph

  • normalize – Apply symmetric normalization

  • rngs (Rngs) – Random number generators

Graph Convolutional Network model with configurable layers and normalization.

Example:

from jraphx.nn.models import GCN
import flax.nnx as nnx

model = GCN(
    in_features=16,
    hidden_features=64,
    num_layers=3,
    out_features=10,
    dropout_rate=0.5,
    norm="layer_norm",  # Options: "batch_norm", "layer_norm", "graph_norm", None
    jk=None,  # Options: "cat", "max", "lstm", None
    rngs=nnx.Rngs(0)
)

# Forward pass
out = model(x, edge_index, batch=batch)
supports_edge_weight: bool = True
supports_edge_attr: bool = False
precompute_norm(edge_index: Array, edge_weight: Array | None = None, num_nodes: int | None = None, dtype: str | type | dtype | None = None) → None[source]

Fill the normalization cache of every GCNConv layer.

Only meaningful for a model built with cached=True. Must be called eagerly, i.e. outside of any JAX transformation, since it mutates module state. Afterwards the model can be called under jax.jit or nnx.jit and reuses the stored normalization.

Parameters:

  • edge_index (Array) – Edge indices [2, num_edges]

  • edge_weight (Array | None, default: None) – Optional edge weights [num_edges]

  • num_nodes (int | None, default: None) – Number of nodes, defaults to the largest node index in edge_index plus one

  • dtype (str | type | dtype | None, default: None) – Data type of the normalized edge weights

Raises:

ValueError – If the model was not built with cached=True.

Return type:

None

init_conv(in_features: int, out_features: int, rngs: Rngs, **kwargs: Any) → MessagePassing[source]

Initialize GCNConv layer.

Every keyword argument reaches the GCNConv constructor, so an unsupported one raises a TypeError instead of being dropped.

Return type:

MessagePassing

GAT

class GAT(*args: Any, **kwargs: Any)[source]

Bases: BasicGNN

Graph Attention Network.

From “Graph Attention Networks” https://arxiv.org/abs/1710.10903 or “How Attentive are Graph Attention Networks?” https://arxiv.org/abs/2105.14491

Uses GATConv or GATv2Conv layers for message passing.

Parameters:

  • in_features (int) – Size of input features

  • hidden_features (int) – Size of hidden layers (per head if concat=True)

  • num_layers (int) – Number of GAT layers

  • out_features (int | None, default: None) – Size of output (if None, uses hidden_features)

  • heads (int, default: 1) – Number of attention heads

  • concat (bool, default: True) – Whether to concatenate or average multi-head outputs

  • v2 (bool, default: False) – Use GATv2Conv instead of GATConv

  • dropout_rate (float, default: 0.0) – Dropout probability

  • act (Callable | None, default: <jax._src.custom_derivatives.custom_jvp object at 0x7f10e9a36ba0>) – Non-linear activation function, or None to disable the activation entirely (default: jax.nn.relu)

  • act_first (bool, default: False) – If True, apply activation before normalization

  • norm (str | None, default: None) – Normalization type (‘batch_norm’, ‘layer_norm’, ‘graph_norm’, None)

  • jk (str | None, default: None) – Jumping Knowledge mode (‘last’, ‘cat’, ‘max’, ‘lstm’, None)

  • residual (bool, default: False) – Whether to use residual connections

  • edge_dim (int | None, default: None) – Edge feature dimension

  • rngs (Rngs) – Random number generators

Graph Attention Network model with multi-head attention and configurable architecture.

Example:

from jraphx.nn.models import GAT
import flax.nnx as nnx

model = GAT(
    in_features=16,
    hidden_features=64,
    num_layers=3,
    out_features=10,
    heads=8,
    v2=False,  # Use GATv2 if True
    dropout_rate=0.6,
    norm="layer_norm",
    jk="max",  # JumpingKnowledge aggregation
    rngs=nnx.Rngs(0)
)

out = model(x, edge_index, batch=batch)
supports_edge_weight: bool = False
supports_edge_attr: bool = True
init_conv(in_features: int, out_features: int, rngs: Rngs, **kwargs: Any) → MessagePassing[source]

Initialize GATConv or GATv2Conv layer.

Remaining keyword arguments (negative_slope, add_self_loops, fill_value, …) reach the convolution constructor, so an unsupported one raises a TypeError instead of being dropped.

Return type:

MessagePassing

GraphSAGE

class GraphSAGE(*args: Any, **kwargs: Any)[source]

Bases: BasicGNN

GraphSAGE: Inductive Representation Learning on Large Graphs.

From “Inductive Representation Learning on Large Graphs” https://arxiv.org/abs/1706.02216

Uses SAGEConv layers for message passing.

Parameters:

  • in_features (int) – Size of input features

  • hidden_features (int) – Size of hidden layers

  • num_layers (int) – Number of GraphSAGE layers

  • out_features (int | None, default: None) – Size of output (if None, uses hidden_features)

  • aggr – Aggregation method (‘mean’, ‘max’, ‘gcn’)

  • dropout_rate (float, default: 0.0) – Dropout probability

  • act (Callable | None, default: <jax._src.custom_derivatives.custom_jvp object at 0x7f10e9a36ba0>) – Non-linear activation function, or None to disable the activation entirely (default: jax.nn.relu)

  • act_first (bool, default: False) – If True, apply activation before normalization

  • norm (str | None, default: None) – Normalization type (‘batch_norm’, ‘layer_norm’, ‘graph_norm’, None)

  • jk (str | None, default: None) – Jumping Knowledge mode (‘last’, ‘cat’, ‘max’, ‘lstm’, None)

  • residual (bool, default: False) – Whether to use residual connections

  • normalize – Whether to L2-normalize output features

  • rngs (Rngs) – Random number generators

GraphSAGE model with multiple aggregation options.

Example:

from jraphx.nn.models import GraphSAGE
import flax.nnx as nnx

model = GraphSAGE(
    in_features=16,
    hidden_features=64,
    num_layers=3,
    out_features=10,
    aggr="mean",  # Options: "mean", "max", "lstm"
    dropout_rate=0.5,
    norm="batch_norm",
    jk="cat",  # Concatenate all layer outputs
    rngs=nnx.Rngs(0)
)

out = model(x, edge_index, batch=batch)
supports_edge_weight: bool = False
supports_edge_attr: bool = False
init_conv(in_features: int, out_features: int, rngs: Rngs, **kwargs: Any) → MessagePassing[source]

Initialize SAGEConv layer.

Every keyword argument reaches the SAGEConv constructor, so an unsupported one raises a TypeError instead of being dropped.

Return type:

MessagePassing

GIN

class GIN(*args: Any, **kwargs: Any)[source]

Bases: BasicGNN

Graph Isomorphism Network.

From “How Powerful are Graph Neural Networks?” https://arxiv.org/abs/1810.00826

Uses GINConv layers with MLP aggregation for message passing.

Parameters:

  • in_features (int) – Size of input features

  • hidden_features (int) – Size of hidden layers

  • num_layers (int) – Number of GIN layers

  • out_features (int | None, default: None) – Size of output (if None, uses hidden_features)

  • dropout_rate (float, default: 0.0) – Dropout probability

  • act (Callable | None, default: <jax._src.custom_derivatives.custom_jvp object at 0x7f10e9a36ba0>) – Non-linear activation function, or None to disable the activation entirely (default: jax.nn.relu)

  • act_first (bool, default: False) – If True, apply activation before normalization

  • norm (str | None, default: None) – Normalization type (‘batch_norm’, ‘layer_norm’, ‘graph_norm’, None). ‘graph_norm’ is applied between GIN blocks; the MLP inside each GINConv uses ‘layer_norm’ instead, because GINConv does not plumb a batch vector into its MLP.

  • jk (str | None, default: None) – Jumping Knowledge mode (‘last’, ‘cat’, ‘max’, ‘lstm’, None)

  • residual (bool, default: False) – Whether to use residual connections

  • train_eps – Whether to learn the epsilon parameter

  • rngs (Rngs) – Random number generators

Graph Isomorphism Network model with MLP transformations.

Example:

from jraphx.nn.models import GIN
import flax.nnx as nnx

model = GIN(
    in_features=16,
    hidden_features=64,
    num_layers=5,
    out_features=10,
    dropout_rate=0.5,
    norm="batch_norm",
    jk="cat",
    rngs=nnx.Rngs(0)
)

out = model(x, edge_index, batch=batch)
supports_edge_weight: bool = False
supports_edge_attr: bool = False
init_conv(in_features: int, out_features: int, rngs: Rngs, **kwargs: Any) → MessagePassing[source]

Initialize GINConv layer.

Remaining keyword arguments (eps, train_eps) reach the GINConv constructor, so an unsupported one raises a TypeError instead of being dropped.

Return type:

MessagePassing

Base Classes

BasicGNN

class BasicGNN(*args: Any, **kwargs: Any)[source]

Bases: Module

An abstract class for implementing basic GNN models.

Subclasses declare which optional edge information their convolution accepts through the supports_edge_weight and supports_edge_attr class attributes, both of which default to False. The forward pass only forwards an argument that the underlying convolution actually consumes, and raises a ValueError when it is handed edge information the subclass has not declared, so that a subclass missing those attributes fails loudly instead of silently training on an unweighted graph.

Parameters:

  • in_features (int) – Size of each input sample, or -1 to derive the size from the first input(s) to the forward method. A tuple corresponds to the sizes of source and target dimensionalities.

  • hidden_features (int) – Size of each hidden sample.

  • num_layers (int) – Number of message passing layers.

  • out_features (int | None, default: None) – If not set to None, will apply a final linear transformation to convert hidden node embeddings to output size out_features. (default: None)

  • dropout_rate (float, default: 0.0) – Dropout probability. (default: 0.)

  • act (Callable | None, default: <jax._src.custom_derivatives.custom_jvp object at 0x7f10e9a36ba0>) – The non-linear activation function to use, or None to disable the activation entirely. (default: jax.nn.relu)

  • act_first (bool, default: False) – If set to True, activation is applied before normalization. (default: False)

  • norm (str | None, default: None) – The normalization function to use ("batch_norm", "layer_norm", "graph_norm" or None). Any other value raises a ValueError. (default: None)

  • jk (str | None, default: None) – The Jumping Knowledge mode ("last", "cat", "max", "lstm"). (default: None)

  • residual (bool, default: False) – Whether to use residual connections between layers. (default: False)

  • rngs (Rngs) – Random number generators for initialization.

  • **kwargs (Any) – Additional arguments for the specific convolution layer.

Abstract base class for GNN models. Provides a common interface for building multi-layer GNNs with normalization, dropout, and JumpingKnowledge connections.

Subclassing Example:

from jraphx.nn.models import BasicGNN
from jraphx.nn.conv import MessagePassing

class MyCustomGNN(BasicGNN):
    def init_conv(self, in_features, out_features, rngs=None, **kwargs):
        # Return your custom message passing layer
        return MyCustomConv(in_features, out_features, rngs=rngs, **kwargs)
supports_edge_weight: bool = False
supports_edge_attr: bool = False
convs: List[MessagePassing]
norms: List[BatchNorm | LayerNorm | GraphNorm | None]
jk: JumpingKnowledge | None
lin: Linear | None
init_conv(in_features: int, out_features: int, rngs: Rngs, **kwargs: Any) → MessagePassing[source]

Initialize convolution layer. To be implemented by subclasses.

Return type:

MessagePassing

Utility Models

MLP

class MLP(*args: Any, **kwargs: Any)[source]

Bases: Module

A Multi-Layer Perception (MLP) model.

There exists two ways to instantiate an MLP:

  1. By specifying explicit feature sizes, e.g.,

    mlp = MLP([16, 32, 64, 128], rngs=nnx.Rngs(0))

    creates a three-layer MLP with differently sized hidden layers.

  2. By specifying fixed hidden feature sizes over a number of layers, e.g.,

    mlp = MLP(in_features=16, hidden_features=32,
          out_features=128, num_layers=3, rngs=nnx.Rngs(0))

    creates a three-layer MLP with equally sized hidden layers.

Parameters:

  • feature_list (Union[list[int], int, None], default: None) – List of input, intermediate and output features such that len(feature_list) - 1 denotes the number of layers of the MLP (default: None)

  • in_features (int | None, default: None) – Size of each input sample. Will override feature_list. (default: None)

  • hidden_features (int | None, default: None) – Size of each hidden sample. Will override feature_list. (default: None)

  • out_features (int | None, default: None) – Size of each output sample. Will override feature_list. (default: None)

  • num_layers (int | None, default: None) – The number of layers. Will override feature_list. (default: None)

  • dropout_rate (float, default: 0.0) – Dropout probability of each hidden embedding. (default: 0.)

  • act (Callable | None, default: <jax._src.custom_derivatives.custom_jvp object at 0x7f10e9a36ba0>) – The non-linear activation function to use, or None to disable the activation entirely. (default: jax.nn.relu)

  • act_first (bool, default: False) – If set to True, activation is applied before normalization. (default: False)

  • norm (str | None, default: None) – The normalization function to use ("batch_norm", "layer_norm" or None). Any other value raises a ValueError. (default: None)

  • plain_last (bool, default: True) – If set to False, will apply non-linearity, batch normalization and dropout to the last layer as well. (default: True)

  • bias (bool, default: True) – If set to False, the module will not learn additive biases. (default: True)

  • rngs (Rngs) – Random number generators for initialization.

Multi-layer perceptron with configurable layers, normalization, and dropout.

Example:

from jraphx.nn.models import MLP
import flax.nnx as nnx

# Using an explicit feature list
mlp = MLP(
    feature_list=[16, 64, 64, 32, 10],
    norm="layer_norm",
    bias=True,
    dropout_rate=0.5,
    act=nnx.relu,
    rngs=nnx.Rngs(0)
)

# Or using in/hidden/out channels
mlp = MLP(
    in_features=16,
    hidden_features=64,
    out_features=10,
    num_layers=3,
    norm="batch_norm",
    dropout_rate=0.5,
    rngs=nnx.Rngs(0)
)

out = mlp(x)
property in_features: int

Size of each input sample.

property out_features: int

Size of each output sample.

property num_layers: int

Number of layers.

JumpingKnowledge

class JumpingKnowledge(*args: Any, **kwargs: Any)[source]

Bases: Module

The Jumping Knowledge layer aggregation module from the “Representation Learning on Graphs with Jumping Knowledge Networks” paper.

Jumping knowledge is performed based on either concatenation ("cat")

\[\mathbf{x}_v^{(1)} \, \Vert \, \ldots \, \Vert \, \mathbf{x}_v^{(T)},\]

max pooling ("max")

\[\max \left( \mathbf{x}_v^{(1)}, \ldots, \mathbf{x}_v^{(T)} \right),\]

or weighted summation

\[\sum_{t=1}^T \alpha_v^{(t)} \mathbf{x}_v^{(t)}\]

with attention scores \(\alpha_v^{(t)}\) obtained from a bi-directional recurrent network ("lstm"). The mode keeps PyG’s name, but the recurrence is a pair of flax.nnx.GRUCell modules run in opposite directions, not an LSTM, so its parameters and outputs differ from PyG’s.

Parameters:

  • mode (str) – The aggregation scheme to use ("cat", "max" or "lstm").

  • num_features (int | None, default: None) – The number of features per representation. Needs to be only set for "lstm"-mode aggregation. (default: None)

  • num_layers (int | None, default: None) – The number of layers to aggregate. Needs to be only set for "lstm"-mode aggregation. (default: None)

  • rngs (Rngs | None, default: None) – Random number generators for initialization.

JumpingKnowledge layer for aggregating representations from different GNN layers.

Example:

from jraphx.nn.models import JumpingKnowledge
import flax.nnx as nnx

# Concatenation mode
jk = JumpingKnowledge(mode="cat", num_features=64, num_layers=3)

# Max pooling mode
jk = JumpingKnowledge(mode="max")

# LSTM aggregation mode
jk = JumpingKnowledge(
    mode="lstm",
    num_features=64,
    num_layers=3,
    rngs=nnx.Rngs(0)
)

# Aggregate layer outputs
layer_outputs = [layer1_out, layer2_out, layer3_out]
final_out = jk(layer_outputs)

Model Selection Guide

Choosing the Right Model

GCN: Best for citation networks and semi-supervised learning tasks with homophilic graphs.

GAT: Excellent for graphs where edge importance varies. The attention mechanism learns which neighbors are most relevant.

GraphSAGE: Ideal for large-scale graphs and inductive learning scenarios where you need to generalize to unseen nodes.

GIN: Most expressive for distinguishing graph structures. Best for graph-level tasks like molecular property prediction.

Configuration Tips

Number of Layers:

  • 2-3 layers for most node classification tasks

  • 4-5 layers for graph-level tasks

  • Use JumpingKnowledge for deeper networks

Normalization:

  • batch_norm: Best for large batches and stable training

  • layer_norm: Works well with smaller batches

  • graph_norm: Specifically designed for graph data

JumpingKnowledge:

  • cat: Preserves all information but increases dimensionality

  • max: Good balance of expressiveness and efficiency

  • lstm: Most flexible but requires more parameters

Dropout:

  • 0.5-0.6 for training stability

  • Higher rates (0.6-0.8) for GAT models

  • Lower rates (0.2-0.5) for deeper models

Performance Comparison

Model Performance Characteristics

Model

Speed

Memory

Expressiveness

Best For

GCN

Fast

Low

Medium

Node classification

GAT

Medium

Medium-High

High

Heterophilic graphs

GraphSAGE

Fast

Low-Medium

Medium

Large-scale graphs

GIN

Fast

Low

Highest

Graph classification