### Install All Jraphx Example Dependencies

Source: https://github.com/dbraun/jraphx/blob/main/examples/README.md

Install all dependencies required for Jraphx examples from the root of the jraphx project. This command ensures all example-specific packages are installed.

```bash
pip install -e ".[examples]"
```

--------------------------------

### Install Development Dependencies

Source: https://github.com/dbraun/jraphx/blob/main/CLAUDE.md

Install the package with development dependencies and set up pre-commit hooks.

```bash
make install-dev
```

--------------------------------

### Run Basic Examples

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/tutorial/examples.md

Execute the Python script containing basic JraphX examples. Ensure you have copied the code into a Python file or Jupyter notebook first.

```bash
python basic_examples.py
```

--------------------------------

### Install in Development Mode

Source: https://github.com/dbraun/jraphx/blob/main/CLAUDE.md

Install the package in editable mode with development dependencies.

```bash
pip install -e ".[dev]"
```

--------------------------------

### Development Installation from Source

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/install/installation.md

Clone the JraphX repository and install it in development mode for the latest features or contributions.

```bash
git clone https://github.com/DBraun/jraphx.git
cd jraphx
pip install -e .
```

--------------------------------

### Install Sphinx Read the Docs Theme

Source: https://github.com/dbraun/jraphx/blob/main/docs/CLAUDE.md

Install the Sphinx Read the Docs theme using pip. This theme is commonly used for project documentation.

```bash
pip install sphinx-rtd-theme
```

--------------------------------

### Generate HTML Documentation

Source: https://github.com/dbraun/jraphx/blob/main/docs/CLAUDE.md

Navigate to the docs directory and run the make html command to generate the documentation. This requires Sphinx to be installed.

```bash
cd docs
make html
```

--------------------------------

### Install JAX AI Stack and JraphX

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/install/installation.md

Install the JAX AI Stack first, then add JraphX. This method ensures compatibility with other JAX ecosystem libraries.

```bash
pip install jax-ai-stack
pip install jraphx
```

--------------------------------

### Development Installation with JAX AI Stack

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/install/installation.md

Combine the JAX AI Stack installation with a development installation of JraphX from its source repository.

```bash
pip install jax-ai-stack
git clone https://github.com/DBraun/jraphx.git
cd jraphx
pip install -e .
```

--------------------------------

### MLP Instantiation with Configuration

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/modules/models.md

Example of creating an MLP instance with specified channel list, normalization, bias, dropout rate, and activation function.

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

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

--------------------------------

### Node Classification Example with Synthetic Data

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/tutorial/examples.md

Provides a complete example for node classification, including synthetic data generation, model initialization, training, and evaluation using JraphX and Optax. This is a good starting point for graph-based classification problems.

```python
import jax
import optax
from jraphx.data import Data

# Create synthetic data
def create_synthetic_data(num_nodes=100, num_features=16, num_classes=4):
    # Use modern Flax NNX Rngs shorthand methods
    rngs = nnx.Rngs(42)

    # Random features
    x = rngs.normal((num_nodes, num_features))

    # Random edges (Erdős-Rényi graph)
    prob = 0.1
    adj = rngs.bernoulli(prob, (num_nodes, num_nodes))
    edge_index = jnp.array(jnp.where(adj)).astype(jnp.int32)

    # Random labels
    y = rngs.randint((num_nodes,), 0, num_classes)

    # Train/val/test splits using indices (JIT-friendly)
    indices = rngs.permutation(jnp.arange(num_nodes))
    train_size = int(0.6 * num_nodes)
    val_size = int(0.8 * num_nodes)

    train_indices = indices[:train_size]
    val_indices = indices[train_size:val_size]
    test_indices = indices[val_size:]

    # Create basic data object
    data = Data(x=x, edge_index=edge_index, y=y)
    return data, train_indices, val_indices, test_indices

# Create data
data, train_indices, val_indices, test_indices = create_synthetic_data()

# Initialize model and optimizer
model = BasicGCN(16, 32, 4, rngs=nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.adam(0.01), wrt=nnx.Param)

# Training function
@nnx.jit
def train_step(model, optimizer, data, train_indices):
    # Ensure model is in training mode
    model.train()

    def loss_fn(model):
        logits = model(data.x, data.edge_index)
        loss = optax.softmax_cross_entropy_with_integer_labels(
            logits[train_indices],
            data.y[train_indices]
        ).mean()
        return loss

    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(model, grads)
    return loss

# Evaluation function
@nnx.jit
def evaluate(model, data, indices):
    # Create evaluation model that shares weights
    eval_model = nnx.merge(*nnx.split(model))
    eval_model.eval()

    logits = eval_model(data.x, data.edge_index)
    preds = jnp.argmax(logits, axis=-1)
    accuracy = (preds[indices] == data.y[indices]).mean()
    return accuracy

```

--------------------------------

### Install PyTorch Geometric Dependencies

Source: https://github.com/dbraun/jraphx/blob/main/examples/README.md

Install PyTorch Geometric for use with Jraphx examples. This command installs the necessary packages for PyTorch Geometric integration.

```bash
pip install torch-geometric
```

--------------------------------

### Multi-Graph Training Example

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/tutorial/gnn_design.md

Demonstrates efficient multi-graph training using Jraphx, Flax NNX, and JAX primitives. Includes graph creation, batching, and a training loop.

```python
import jax
import jax.numpy as jnp
from flax import nnx
from jraphx.data import Data, Batch
from jraphx.nn.pool import global_mean_pool
import optax

# Dummy GCN model for demonstration
class GCN(nnx.Module):
    def __init__(self, in_features, hidden_features, out_features, *, rngs: nnx.Rngs):
        self.layer1 = nnx.Linear(in_features, hidden_features, rngs=rngs)
        self.layer2 = nnx.Linear(hidden_features, out_features, rngs=rngs)

    def __call__(self, x):
        x = self.layer1(x)
        x = jax.nn.relu(x)
        x = self.layer2(x)
        return x

# Create multiple training graphs using new Rngs shorthand methods
rngs = nnx.Rngs(0, params=1)  # Separate keys for different purposes
train_graphs = []
num_nodes = 0
for i in range(100):
    # Use Rngs shorthand methods (Flax NNX feature)
    n_nodes = rngs.randint((), 10, 50)  # Much cleaner than random.randint!
    x = rngs.params.normal((n_nodes, 16))  # Use params key for features
    # Create random edges (simplified)
    n_edges = n_nodes - 1
    edge_index = jnp.stack([
        jnp.arange(n_edges),
        jnp.roll(jnp.arange(n_edges), 1)
    ])
    train_graphs.append(Data(x=x, edge_index=edge_index))

# Batch training function
@nnx.jit
def train_on_batch(model, optimizer, graphs, targets):
    batch = Batch.from_data_list(graphs)

    def loss_fn(model):
        predictions = model(batch.x) # Assuming model takes node features directly
        # Global pooling to get graph-level predictions
        graph_preds = global_mean_pool(predictions, batch.batch)
        return jnp.mean((graph_preds - targets) ** 2)

    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(model, grads)
    return loss

# Training loop
model_rngs = nnx.Rngs(42)  # For model initialization
model = GCN(16, 32, 7, rngs=model_rngs)
optimizer = nnx.Optimizer(model, optax.adam(0.01), wrt=nnx.Param)

target_rngs = nnx.Rngs(100)  # Separate Rngs for targets
for epoch in range(50):
    # Sample batch of graphs
    batch_graphs = train_graphs[:32]  # Batch size 32
    batch_targets = target_rngs.normal((32, 7))  # Shorthand method!

    loss = train_on_batch(model, optimizer, batch_graphs, batch_targets)
    if epoch % 10 == 0:
        print(f'Epoch {epoch}, Loss: {loss:.4f}')

```

--------------------------------

### Setup Optimizer and Training Loop

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/get_started/introduction.md

Initializes an Adam optimizer for the GCN model and defines a JIT-compiled training step function. The loop trains the model for a specified number of epochs.

```python
from flax import nnx
import optax

# Setup optimizer
optimizer = nnx.Optimizer(model, optax.adam(0.01), wrt=nnx.Param)

# Training loop
train_mask = jnp.array([True, True, False, False])  # First 2 nodes for training
test_mask = jnp.array([False, False, True, True])   # Last 2 nodes for testing

@nnx.jit
def train_step(model, optimizer, data, train_mask):
    def loss_fn_inner(model):
        return loss_fn(model, data, train_mask)

    loss, grads = nnx.value_and_grad(loss_fn_inner)(model)
    optimizer.update(model, grads)
    return loss

# Train for a few epochs
model.train()
for epoch in range(200):
    loss = train_step(model, optimizer, data, train_mask)
    if epoch % 50 == 0:
        print(f'Epoch {epoch}, Loss: {loss:.4f}')
```

--------------------------------

### Understanding JAX Batch Output

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/get_started/introduction.md

This example shows the typical output structure when printing a JAX batch object, indicating the sizes of different graph components.

```python
print(batch)
>>> Batch(batch=[1082], edge_index=[2, 4066], x=[1082, 21], y=[32])
```

--------------------------------

### Create Example Graph Data

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/get_started/introduction.md

Initializes a graph with node features, edge indices, and node labels for use in GNN models. Ensure `jraphx.data.Data` and `jax.numpy` are imported.

```python
import jax.numpy as jnp
from jraphx.data import Data

# Create a simple graph with 4 nodes, 3 features per node, 3 classes
x = jnp.array([[1.0, 0.5, 0.2], [0.8, 1.0, 0.1], [0.3, 0.7, 1.0], [0.9, 0.2, 0.8]], dtype=jnp.float32)
edge_index = jnp.array([[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]], dtype=jnp.int32)  # Undirected edges
y = jnp.array([0, 0, 1, 1], dtype=jnp.int32)  # Node labels

data = Data(x=x, edge_index=edge_index, y=y)
print(f"Graph: {data.num_nodes} nodes, {data.num_edges} edges")
```

--------------------------------

### Install JraphX from PyPI

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/install/installation.md

Use this command to install the latest stable version of JraphX and its dependencies from the Python Package Index.

```bash
pip install jraphx
```

--------------------------------

### Generate PyTorch Geometric Installation Command

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/install/quick-start.html

This snippet dynamically generates the installation command for PyTorch Geometric based on selected options for PyTorch version, OS, package manager, and CUDA version. It also displays compatibility warnings for unsupported combinations.

```javascript
var torchList = [
  ['torch-2.8.0', 'PyTorch 2.8.*'],
  ['torch-2.7.0', 'PyTorch 2.7.*'],
  ['torch-2.6.0', 'PyTorch 2.6.*'],
];
var osList = [
  ['linux', 'Linux'],
  ['mac', 'Mac'],
  ['windows', 'Windows'],
];
var packageList = [
  ['pip', 'Pip'],
  ['conda', 'Conda'],
];
var cudaList = [
  ['cu118', '11.8'],
  ['cu121', '12.1'],
  ['cu124', '12.4'],
  ['cu126', '12.6'],
  ['cu128', '12.8'],
  ['cu129', '12.9'],
  ['cpu', 'CPU'],
];
torchList.forEach(x => $("#torch").append(\`<div id="${x[0]}">${x[1]}</div>\`));
osList.forEach(x => $("#os").append(\`<div id="${x[0]}">${x[1]}</div>\`));
packageList.forEach(x => $("#package").append(\`<div id="${x[0]}">${x[1]}</div>\`));
cudaList.forEach(x => $("#cuda").append(\`<div id="${x[0]}">${x[1]}</div>\`));

function updateCommand() {
  var torch = $("#command").attr("torch");
  var os = $("#command").attr("os");
  var package = $("#command").attr("package");
  var cuda = $("#command").attr("cuda");

  if (os == "mac" && cuda != "cpu") {
    $("#command pre").text('# macOS binaries do not support CUDA');
  } else if (torch == "torch-2.6.0" && cuda == "cu121") {
    $("#command pre").text('# PyTorch version does not support CUDA 12.1');
  } else if (torch == "torch-2.6.0" && cuda == "cu128") {
    $("#command pre").text('# PyTorch version does not support CUDA 12.8');
  } else if (torch == "torch-2.7.0" && cuda == "cu121") {
    $("#command pre").text('# PyTorch version does not support CUDA 12.1');
  } else if (torch == "torch-2.7.0" && cuda == "cu124") {
    $("#command pre").text('# PyTorch version does not support CUDA 12.4');
  } else if (torch == "torch-2.8.0" && cuda == "cu118") {
    $("#command pre").text('# PyTorch version does not support CUDA 11.8');
  } else if (torch == "torch-2.8.0" && cuda == "cu121") {
    $("#command pre").text('# PyTorch version does not support CUDA 11.8');
  } else if (torch == "torch-2.8.0" && cuda == "cu124") {
    $("#command pre").text('# PyTorch version does not support CUDA 12.4');
  } else if (package == "conda") {
    $("#command pre").text('# Conda packages are no longer available since PyTorch >2.5.0. Please use pip instead.');
  } else {
    $("#command pre").text(\`pip install torch_geometric\n\n# Optional dependencies:\npip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/${$("#command").attr("torch")}+${$("#command").attr("cuda")}.html
ylic`);
  }
}

$(".quick-start .content-column .row div").click(function() {
  $(this).parent().children().removeClass("selected");
  $(this).addClass("selected");
  $("#command").attr($(this).parent().attr("id"), $(this).attr("id"));
  updateCommand();
});

$("#torch").children().get(0).click();
$("#linux").click();
$("#pip").click();
$("#cpu").click();
```

--------------------------------

### Node Degree Calculation Example

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/modules/utils.md

Demonstrates calculating node degrees using a simple index array. The `num_nodes` parameter can be used to specify the total number of nodes.

```python
import jax.numpy as jnp
from jraphx.utils import degree

row = jnp.array([0, 1, 0, 2, 0])
print(degree(row, dtype=jnp.int32))
```

--------------------------------

### Serve Documentation Locally

Source: https://github.com/dbraun/jraphx/blob/main/docs/CLAUDE.md

Launch a local HTTP server to view the generated documentation. The -d build/html flag specifies the directory containing the HTML files.

```bash
python -m http.server -d build/html
```

--------------------------------

### JIT-Compile a Training Step

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/advanced/jit.md

Shows how to JIT compile an entire training step for optimal performance. It includes loss calculation, gradient computation, and optimizer updates. Uses concrete indices for JIT compatibility.

```python
import optax

# Setup optimizer
optimizer = nnx.Optimizer(model, optax.adam(0.01), wrt=nnx.Param)

@jax.jit
def train_step(model, optimizer, x, edge_index, targets, train_indices):
    """JIT-compiled training step."""
    def loss_fn(model):
        predictions = model(x, edge_index)
        # Use concrete indices instead of boolean mask for JIT compatibility
        train_predictions = predictions[train_indices]
        train_targets = targets[train_indices]
        return jnp.mean(optax.softmax_cross_entropy_with_integer_labels(
            train_predictions, train_targets
        ))

    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(model, grads)
    return loss

# Training loop with JIT compilation
targets = jnp.array([0, 1, 2, 0, 1, 2, 0] * 14 + [0, 1, 2])  # 100 targets
train_indices = jnp.arange(80)  # First 80 nodes for training (concrete indices)

for epoch in range(100):
    loss = train_step(model, optimizer, data.x, data.edge_index, targets, train_indices)
    if epoch % 20 == 0:
        print(f'Epoch {epoch}, Loss: {loss:.4f}')
```

--------------------------------

### Verify JraphX Installation

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/install/installation.md

Run this Python script to check installed versions of JAX, JraphX, and test basic graph processing functionality.

```python
import jax
import jax.numpy as jnp
from flax import nnx
import jraphx

print(f"JAX version: {jax.__version__}")
print(f"JAX backend: {jax.default_backend()}")
print(f"JraphX version: {jraphx.__version__}")

# Test basic functionality
from jraphx.data import Data
from jraphx.nn.conv import GCNConv

# Create a simple graph
data = Data(
    x=jnp.ones((3, 4)),
    edge_index=jnp.array([[0, 1, 2], [1, 2, 0]])
)

# Create and use a GNN layer
layer = GCNConv(4, 8, rngs=nnx.Rngs(42))
output = layer(data.x, data.edge_index)
print(f"Successfully processed graph: {output.shape}")
```

--------------------------------

### Initializing and Calling GCNConv Layer

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/notes/create_gnn.md

Demonstrates how to instantiate and use the GCNConv layer with sample node features and edge indices.

```python
conv = GCNConv(16, 32, rngs=nnx.Rngs(42))
output = conv(x, edge_index)
```

--------------------------------

### Basic Graph Batching with JraphX

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/advanced/batching.md

Demonstrates how to create individual graphs and then batch them together using the `jraphx.data.Batch` class. Shows how to inspect the properties of the created batch.

```python
import jax.numpy as jnp
from jraphx.data import Data, Batch

# Create individual graphs
graph1 = Data(
    x=jnp.array([[1.0, 2.0], [3.0, 4.0]]),  # 2 nodes
    edge_index=jnp.array([[0, 1], [1, 0]])   # 2 edges
)

graph2 = Data(
    x=jnp.array([[5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]),  # 3 nodes
    edge_index=jnp.array([[0, 1, 2], [1, 2, 0]])          # 3 edges
)

# Create batch
batch = Batch.from_data_list([graph1, graph2])

print(f"Batch info:")
print(f"  Total nodes: {batch.num_nodes}")      # 5 nodes total
print(f"  Total edges: {batch.num_edges}")      # 5 edges total
print(f"  Num graphs: {batch.num_graphs}")      # 2 graphs
print(f"  Node features: {batch.x.shape}")     # [5, 2]
print(f"  Edge indices: {batch.edge_index.shape}")  # [2, 5]
print(f"  Batch vector: {batch.batch}")        # [0, 0, 1, 1, 1]
```

--------------------------------

### Basic Graph Preprocessing

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/notes/introduction.md

Applies a default preprocessing function to graph data. This is a starting point for data preparation.

```python
import jraphx as jraphx
from jraphx.data import Data

original_data = Data(x=jnp.ones((3, 2)), edge_index=jnp.array([[0, 1], [1, 2]]))
processed_data = preprocess_graph(original_data)
```

--------------------------------

### JAX JIT Compilation and Vectorization for GNNs

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/cheatsheet/gnn_cheatsheet.md

Illustrates how to leverage JAX's `@jax.jit` for JIT compilation and `nnx.vmap` for vectorizing GNN inference over multiple graphs. Also shows integration with Optax for optimization.

```python
import jax

# JIT compile for speed
@jax.jit
def fast_gnn_inference(model, x, edge_index):
    return model(x, edge_index)

# Vectorize over multiple graphs (fixed-size)
@nnx.vmap
def batch_gnn_inference(x_batch, edge_index_batch):
    return model(x_batch, edge_index_batch)

# Use with optimization libraries
import optax
optimizer = nnx.Optimizer(model, optax.adam(0.01), wrt=nnx.Param)

@jax.jit
def train_step(model, optimizer, data, targets):
    def loss_fn(model):
        preds = model(data.x, data.edge_index)
        return jnp.mean((preds - targets) ** 2)

    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(model, grads)
    return loss
```

--------------------------------

### LayerNorm

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/modules/norm.md

Applies layer normalization over each individual example in a batch of node features. Supports both node-wise and graph-wise normalization modes.

```APIDOC
## LayerNorm

### Description
Applies layer normalization over each individual example in a batch of node features.

### Parameters
* **num_features** (*int* *or* *list*) – Size of each input sample, or list of dimensions to normalize.
* **eps** (*float* *,* *optional*) – A value added to the denominator for numerical stability. (default: `1e-5`)
* **elementwise_affine** (*bool* *,* *optional*) – If set to `True`, this module has learnable affine parameters $\gamma$ and $\beta$. (default: `True`)
* **mode** (*str* *,* *optional*) – The normalization mode to use for layer normalization (`"graph"` or `"node"`). (default: `"node"`)
* **dtype** – The dtype of the result (default: infer from input and params).
* **param_dtype** – The dtype passed to parameter initializers (default: float32).
* **use_bias** (*bool* *,* *optional*) – If True, bias (beta) is added. (default: `True`)
* **use_scale** (*bool* *,* *optional*) – If True, multiply by scale (gamma). (default: `True`)
* **bias_init** – Initializer for bias, by default, zero.
* **scale_init** – Initializer for scale, by default, one.
* **reduction_axes** – Axes for computing normalization statistics.
* **feature_axes** – Feature axes for learned bias and scaling.
* **axis_name** – The axis name used to combine batch statistics from multiple devices.
* **axis_index_groups** – Groups of axis indices within that named axis.
* **use_fast_variance** – If true, use faster, but less numerically stable variance calculation.
* **rngs** – Random number generators for initialization.

### Example
```python
from jraphx.nn.norm import LayerNorm
import flax.nnx as nnx

# Node-wise normalization
norm = LayerNorm(
    num_features=64,
    mode="node",
    eps=1e-5,
    elementwise_affine=True,
    rngs=nnx.Rngs(0)
)

x_normalized = norm(x)

# Graph-wise normalization (requires batch index)
norm_graph = LayerNorm(
    num_features=64,
    mode="graph",
    rngs=nnx.Rngs(0)
)

x_normalized = norm_graph(x, batch=batch)
```
```

--------------------------------

### Custom GNN Subclassing

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/modules/models.md

Example of subclassing the abstract BasicGNN class to define a custom GNN model by implementing the init_conv method.

```python
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)
```

--------------------------------

### Simple GNN Training Loop with Optax Optimizer

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/get_started/introduction.md

Implement a basic training step using `nnx.jit` and `nnx.value_and_grad`. This snippet sets up an Adam optimizer and defines a loss function for training.

```python
import optax
from jraphx.data import DataLoader

# Create optimizer
optimizer = nnx.Optimizer(model, optax.adam(learning_rate=0.01), wrt=nnx.Param)

@nnx.jit
def train_step(model, optimizer, data, labels):
    # Ensure model is in training mode
    model.train()

    def loss_fn(model):
        logits = model(data.x, data.edge_index)
        loss = optax.softmax_cross_entropy(logits, labels).mean()
        return loss

    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(model, grads)
    return loss

# Training loop
for epoch in range(100):
    loss = train_step(model, optimizer, data, labels)
    if epoch % 10 == 0:
        print(f"Epoch {epoch}, Loss: {loss:.4f}")
```

--------------------------------

### Custom GNN with GraphNorm Integration

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/modules/norm.md

Example of integrating GraphNorm within a custom GNN module for node feature transformation and normalization.

```python
from jraphx.nn.conv import GCNConv
from jraphx.nn.norm import GraphNorm
import flax.nnx as nnx

class CustomGNN(nnx.Module):
    def __init__(self, in_features, out_features, rngs):
        self.conv = GCNConv(in_features, 64, rngs=rngs)
        self.norm = GraphNorm(64, rngs=rngs)
        self.linear = nnx.Linear(64, out_features, rngs=rngs)

    def __call__(self, x, edge_index, batch=None):
        x = self.conv(x, edge_index)
        x = self.norm(x, batch)
        x = nnx.relu(x)
        return self.linear(x)
```

--------------------------------

### Initialize and Use GCN, GAT, and SAGE Convolution Layers

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/get_started/introduction.md

Demonstrates the initialization and basic usage of GCNConv, GATConv, and SAGEConv layers from JraphX for graph convolution operations.

```python
import flax.nnx as nnx
from jraphx.nn.conv import GCNConv, GATConv, SAGEConv

# Initialize random number generator
rngs = nnx.Rngs(42)

# Graph Convolutional Network (GCN)
gcn = GCNConv(in_features=3, out_features=16, rngs=rngs)
out = gcn(data.x, data.edge_index)

# Graph Attention Network (GAT)
gat = GATConv(in_features=3, out_features=16, heads=4, rngs=rngs)
out = gat(data.x, data.edge_index)

# GraphSAGE
sage = SAGEConv(in_features=3, out_features=16, rngs=rngs)
out = sage(data.x, data.edge_index)
```

--------------------------------

### Run Pre-commit Hooks

Source: https://github.com/dbraun/jraphx/blob/main/CLAUDE.md

Execute all configured pre-commit hooks to verify code quality before committing.

```bash
make pre-commit
```

--------------------------------

### TopKPooling Layer Initialization and Application

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/modules/pooling.md

Initializes and applies the TopKPooling layer to select the top 50% of nodes based on learnable scores. Ensure 'x', 'edge_index', 'edge_attr', and 'batch' are defined before use.

```python
from jraphx.nn.pool import TopKPooling
import flax.nnx as nnx

# Select top 50% of nodes
pool = TopKPooling(
    in_features=64,
    ratio=0.5,
    min_score=None,  # Optional minimum score threshold
    multiplier=1.0,  # Score multiplier
    rngs=nnx.Rngs(0)
)

# Apply pooling
x_pool, edge_index_pool, edge_attr_pool, batch_pool, perm = pool(
    x, edge_index, edge_attr=edge_attr, batch=batch
)

# perm contains indices of selected nodes
```

--------------------------------

### Custom Aggregation in MessagePassing

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/modules/conv.md

Implement custom aggregation logic by subclassing MessagePassing and overriding the aggregate method. This example demonstrates using scatter_mean for custom aggregation.

```python
class CustomConv(MessagePassing):
    def __init__(self, in_features, out_features):
        # Custom aggregation function
        super().__init__(aggr='add')

    def aggregate(self, inputs, index, dim_size=None):
        # Override for custom aggregation
        return scatter_mean(inputs, index, dim=0, dim_size=dim_size)
```

--------------------------------

### JIT Compile Graph Processing Function

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/cheatsheet/data_cheatsheet.md

Compiles a graph processing function using JAX's JIT for performance. This example adds self-loops to a graph.

```python
import jax

@jax.jit
def process_graph(data):
    from jraphx.utils import add_self_loops
    edge_index, _ = add_self_loops(data.edge_index, data.x.shape[0])
    return edge_index

processed = process_graph(data)
```

--------------------------------

### JAX JIT Compilation for Batched Graph Processing

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/advanced/batching.md

Shows how to use JAX's JIT compilation with a GCNConv model to efficiently process batched graph data. Includes a helper function for processing lists of graphs.

```python
import jax
from jraphx.nn.conv import GCNConv
from flax import nnx

# Create model
model = GCNConv(2, 8, rngs=nnx.Rngs(42))

# Process batch with JIT compilation (extract arrays first)
@jax.jit
def process_batch(model, x, edge_index):
    return model(x, edge_index)

# Efficient batch processing
batch_output = process_batch(model, batch.x, batch.edge_index)
print(f"Batch output shape: {batch_output.shape}")  # [5, 8]

# For multiple batches, process arrays directly
def process_graph_list(model, graph_list):
    """Process a list of graphs efficiently."""
    batch = Batch.from_data_list(graph_list)
    return process_batch(model, batch.x, batch.edge_index)
```

--------------------------------

### Batch Processing with Graph-Level Targets

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/advanced/batching.md

Demonstrates how to create graphs with graph-level targets and process them in batches using `global_mean_pool` for aggregation.

```python
from jraphx.nn.pool import global_mean_pool

# Create graphs with graph-level targets
graphs_with_targets = []
rngs = nnx.Rngs(0, targets=1)  # Use Flax NNX shorthand

for i in range(10):
    x = rngs.normal((5, 16))  # Node features
    edge_index = jnp.array([[0, 1, 2], [1, 2, 0]])  # Simple cycle
    target = rngs.targets.normal((7,))  # Graph-level target

    graph = Data(x=x, edge_index=edge_index)
    graphs_with_targets.append((graph, target))

# Batch processing with graph-level targets
def process_graph_batch_with_targets(model, graphs_and_targets):
    graphs, targets = zip(*graphs_and_targets)

    # Create batch for graphs
    batch = Batch.from_data_list(graphs)

    # Process batch
    node_embeddings = model(batch.x, batch.edge_index)

    # Pool to graph-level
    graph_embeddings = global_mean_pool(node_embeddings, batch.batch)

    # Stack targets to create [num_graphs, target_dim]
    targets_array = jnp.stack(targets)

    return graph_embeddings, targets_array
```

--------------------------------

### JraphX Model Compilation with JIT

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/advanced/compile.md

Demonstrates how to create a JraphX model and use JAX's JIT (Just-In-Time) compilation to optimize the forward pass for performance. XLA optimization is triggered automatically by JIT.

```python
import jax
import jax.numpy as jnp
from jraphx.nn.models import GCN
from flax import nnx

# Create model - XLA will optimize this automatically when JIT-compiled
model = GCN(
    in_features=64,
    hidden_features=128,
    out_features=32,
    num_layers=4,
    rngs=nnx.Rngs(42)
)

# JIT compilation triggers XLA optimization
@nnx.jit
def optimized_forward(model, x, edge_index):
    return model(x, edge_index)

# XLA optimizes the entire computation graph
x = jnp.ones((1000, 64))
edge_index = jnp.array([[0, 1, 2], [1, 2, 0]])
output = optimized_forward(model, x, edge_index)
```

--------------------------------

### Memory-Efficient Deep GNN with nnx.scan

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/tutorial/gnn_design.md

Implement a deep GNN using nnx.scan for memory-efficient sequential processing of layers. This example defines a DeepGNN module with multiple hidden blocks.

```python
from jraphx.nn.conv import GCNConv

class HiddenBlock(nnx.Module):
    """Single hidden layer block for scanning."""
    def __init__(self, hidden_features: int, rngs: nnx.Rngs):
        self.conv = GCNConv(hidden_features, hidden_features, rngs=rngs)

    def __call__(self, x, edge_index):
        x = self.conv(x, edge_index)
        x = nnx.relu(x)
        return x  # Return only x, no second output needed

class DeepGNN(nnx.Module):
    def __init__(self, in_features: int, hidden_features: int, out_features: int, num_layers: int, rngs: nnx.Rngs):
        # Create input and output layers
        self.input_layer = GCNConv(in_features, hidden_features, rngs=rngs)
        self.output_layer = GCNConv(hidden_features, out_features, rngs=rngs)

        # Create multiple hidden layers using vmap
        num_hidden = num_layers - 2
        self.num_hidden = num_hidden

        if num_hidden > 0:
            @nnx.split_rngs(splits=num_hidden)
            @nnx.vmap(in_axes=(0,), out_axes=0)
            def create_hidden_block(rngs: nnx.Rngs):
                return HiddenBlock(hidden_features, rngs=rngs)

            self.hidden_blocks = create_hidden_block(rngs)
        else:
            self.hidden_blocks = None

    def __call__(self, data):
        x, edge_index = data.x, data.edge_index

        # Input layer
        x = self.input_layer(x, edge_index)
        x = nnx.relu(x)

        # Hidden layers with scan (only if we have hidden layers)
        if self.num_hidden > 0:
            @nnx.scan(in_axes=(nnx.Carry, 0), out_axes=nnx.Carry)
            def forward_hidden(x, block):
                x = block(x, edge_index)
                return x

            x = forward_hidden(x, self.hidden_blocks)

        # Output layer
        return self.output_layer(x, edge_index)

# Create and use deep network
deep_model = DeepGNN(16, 64, 7, 10, rngs=nnx.Rngs(42))
deep_predictions = deep_model(data)
```

--------------------------------

### Model Initialization: PyTorch Geometric vs JraphX

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/missing_tests.md

Demonstrates the difference in initializing a GCNConv layer between PyTorch Geometric and JraphX, highlighting the need for RNGs in JraphX.

```python
# PyTorch Geometric
conv = GCNConv(16, 32)

# JraphX
conv = GCNConv(16, 32, rngs=nnx.Rngs(0))
```

--------------------------------

### Edge-Conditioned Convolution Implementation

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/advanced/techniques.md

Implement custom message passing layers, such as edge-conditioned convolutions, by subclassing MessagePassing. This example defines how messages are constructed using node and edge features.

```python
from jraphx.nn.conv import MessagePassing
import flax.nnx as nnx

class EdgeConditionedConv(MessagePassing):
    """Message passing with edge features."""

    def __init__(self, in_features, out_features, edge_dim, rngs):
        super().__().__init__(aggr='mean')
        self.node_mlp = nnx.Sequential(
            nnx.Linear(in_features * 2 + edge_dim, out_features, rngs=rngs),
            nnx.relu,
            nnx.Linear(out_features, out_features, rngs=rngs)
        )

    def message(self, x_i, x_j, edge_attr):
        # Concatenate source, target, and edge features
        msg = jnp.concatenate([x_i, x_j, edge_attr], axis=-1)
        return self.node_mlp(msg)

    def __call__(self, x, edge_index, edge_attr):
        return self.propagate(
            edge_index, x=x, edge_attr=edge_attr
        )
```

--------------------------------

### SAGEConv Initialization and Usage

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/modules/conv.md

Initializes and uses the SAGEConv layer with mean and LSTM aggregation. Supports normalization and requires RNGs for initialization.

```python
from jraphx.nn.conv import SAGEConv
import flax.nnx as nnx

# Mean aggregation (most common)
conv = SAGEConv(
    in_features=16,
    out_features=32,
    aggr='mean',
    normalize=True,
    rngs=nnx.Rngs(0)
)

# LSTM aggregation
conv_lstm = SAGEConv(
    in_features=16,
    out_features=32,
    aggr='lstm',
    rngs=nnx.Rngs(0)
)

out = conv(x, edge_index)
```

--------------------------------

### Process Multiple Graphs with JAX vmap

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/get_started/introduction.md

Demonstrates how to prepare for processing multiple graphs in parallel using jax.vmap. This requires proper batching of graph data.

```python
# Create multiple graphs
graphs = [Data(x=jnp.ones((3, 2)), edge_index=jnp.array([[0, 1], [1, 0]]))
          for _ in range(5)]

# Process multiple graphs in parallel
def process_single_graph(data):
    return jnp.sum(data.x)

# vmap over a batch of graphs
batched_process = jax.vmap(process_single_graph)
# results = batched_process(graph_batch)  # Requires proper batching
```

--------------------------------

### Build a Complete GNN Model with JraphX Layers

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/get_started/introduction.md

Combine GCNConv and global_mean_pool layers to construct a full GNN model. This example demonstrates a typical architecture for graph-level prediction tasks.

```python
import jax
import flax.nnx as nnx
from jraphx.nn.conv import GCNConv
from jraphx.nn.pool import global_mean_pool

class GNN(nnx.Module):
    def __init__(self, in_features, hidden_features, out_features, rngs):
        self.conv1 = GCNConv(in_features, hidden_features, rngs=rngs)
        self.conv2 = GCNConv(hidden_features, hidden_features, rngs=rngs)
        self.conv3 = GCNConv(hidden_features, out_features, rngs=rngs)
        self.dropout = nnx.Dropout(rate=0.5, rngs=rngs)

    def __call__(self, x, edge_index, batch=None):
        # First GCN layer
        x = self.conv1(x, edge_index)
        x = nnx.relu(x)
        x = self.dropout(x)

        # Second GCN layer
        x = self.conv2(x, edge_index)
        x = nnx.relu(x)
        x = self.dropout(x)

        # Third GCN layer
        x = self.conv3(x, edge_index)

        # Global pooling (for graph-level prediction)
        if batch is not None:
            x = global_mean_pool(x, batch)

        return x

# Create model
model = GNN(in_features=3, hidden_features=64, out_features=10, rngs=nnx.Rngs(42))

# Forward pass
output = model(data.x, data.edge_index)
```

--------------------------------

### Gradual Pooling for Better Gradients

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/modules/pooling.md

Illustrates a strategy of applying pooling gradually over multiple layers, starting with a higher ratio and decreasing it in subsequent layers. This approach can help maintain gradient flow.

```python
# Gradual pooling for better gradients
pool1 = TopKPooling(64, ratio=0.8, rngs=nnx.Rngs(42))  # First layer: 80%
pool2 = TopKPooling(64, ratio=0.6, rngs=nnx.Rngs(42))  # Second layer: 60%
```

--------------------------------

### Tensor Operations: PyTorch vs JAX

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/missing_tests.md

Illustrates the conversion of tensor creation and comparison operations between PyTorch and JAX.

```python
# PyTorch
x = torch.randn(4, 16)
assert torch.allclose(x, x)

# JAX
x = random.normal(random.key(0), (4, 16))
assert jnp.allclose(x, x)
```

--------------------------------

### Manage Train/Eval Modes for GNNs with Dropout

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/get_started/introduction.md

Set models to training or evaluation mode using `model.train()` and `model.eval()`. This example shows how dropout is enabled during training and disabled during evaluation, while weights remain synchronized.

```python
from jraphx.nn.models import GraphSAGE

# Create model with dropout
model = GraphSAGE(in_features=16, hidden_features=32, out_features=8,
                 num_layers=2, dropout_rate=0.5, rngs=nnx.Rngs(42))
model.train()  # Set to training mode

# Create evaluation model that shares weights
eval_model = nnx.merge(*nnx.split(model))  # Same weights, different behavior
eval_model.eval()  # Set to evaluation mode

# Both models share weights but behave differently
train_out = model(x, edge_index)      # Uses dropout
eval_out = eval_model(x, edge_index)  # No dropout

# Weights stay synchronized automatically - no copying needed!
print("Weights shared:", jnp.allclose(
    model.convs[0].linear.kernel.value,
    eval_model.convs[0].linear.kernel.value
))
>>> Weights shared: True
```

--------------------------------

### Custom vmap Pattern for Graph Aggregation

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/advanced/techniques.md

Implement custom vmap patterns for specific graph processing tasks, such as aggregating features across a batch of graphs. This example defines a per-graph operation and then vectorizes it.

```python
def custom_vmap_aggregation(graphs, model):
    """Custom vmap pattern for graph aggregation."""

    # Define per-graph operation
    def per_graph_op(graph):
        node_features = model(graph.x, graph.edge_index)
        # Custom aggregation
        graph_feature = node_features.mean(axis=0)
        return graph_feature

    # Vectorize and apply
    vmapped_op = nnx.vmap(per_graph_op)
    graph_features = vmapped_op(graphs)

    # Further processing on all graphs
    return graph_features.mean(axis=0)
```

--------------------------------

### Debugging Compiled Code with Intermediate Prints

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/advanced/compile.md

Provides an example of using `jax.debug.print` within a function to inspect intermediate values and shapes during execution. Note that this specific method for printing only works in eager mode.

```python
# Print intermediate values (only works in eager mode)
def debug_forward(model, x, edge_index):
    x = model.layers[0](x, edge_index)
    jax.debug.print("After layer 0: {}", x.shape)
    x = model.layers[1](x, edge_index)
    jax.debug.print("After layer 1: {}", x.shape)
    return x
```

--------------------------------

### Create a JIT-Compatible Custom GNN Layer

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/advanced/jit.md

Illustrates how to create a custom JraphX layer that is JIT-compatible. Emphasizes using only JAX operations, static shapes, and pure functions.

```python
from jraphx.nn.conv import MessagePassing

class CustomGNNLayer(MessagePassing):
    def __init__(self, in_features, out_features, *, rngs: nnx.Rngs):
        super().__init__(aggr='mean')
        self.linear = nnx.Linear(in_features, out_features, rngs=rngs)

    def __call__(self, x, edge_index):
        # All operations here must be JAX-compatible
        x = self.linear(x)

        # Use JAX operations for conditionals
        x = jnp.where(x > 0, x, 0.0)  # ReLU activation

        # Standard message passing
        return self.propagate(edge_index, x)

# This layer is automatically JIT-compatible
@jax.jit
def forward_with_custom_layer(x, edge_index):
    layer = CustomGNNLayer(16, 32, rngs=nnx.Rngs(42))
    return layer(x, edge_index)
```

--------------------------------

### GINConv Initialization with MLP

Source: https://github.com/dbraun/jraphx/blob/main/docs/source/modules/conv.md

Initializes the GINConv layer, which uses a provided MLP for transformation. The epsilon parameter can be fixed or trainable.

```python
from jraphx.nn.conv import GINConv
from jraphx.nn.models import MLP
import flax.nnx as nnx
```