Chapter 3: Message Passing and GNNs

Learning Objectives

By reading this chapter, you will be able to:

• ✅ Understand the basic structure of the message passing framework (Message, Aggregate, Update)
• ✅ Master the mathematical formulation of generalized GNNs (MPNN)
• ✅ Implement GraphSAGE's sampling-based aggregation
• ✅ Understand the characteristics of various aggregators (Mean, Pool, LSTM)
• ✅ Understand the relationship between GIN (Graph Isomorphism Network) and the WL test
• ✅ Evaluate the discriminative (expressive) power of GNNs
• ✅ Master efficient implementation with PyTorch Geometric
• ✅ Implement graph classification tasks and batch processing

3.1 The Message Passing Framework

The Concept of Message Passing

Message passing is a framework that provides a unified description of information propagation in GNNs. Nodes exchange messages with one another and aggregate them to update their features.

"The message passing framework provides a unified way to describe any GNN architecture in terms of three basic operations: Message, Aggregate, and Update."

The Three Basic Operations

Message passing consists of the following three steps:

Step 1: Message (Message Generation)

Generate the messages sent from neighboring nodes to the center node:

$$ \mathbf{m}_{j \to i}^{(k)} = \text{MESSAGE}^{(k)}\left(\mathbf{h}_i^{(k-1)}, \mathbf{h}_j^{(k-1)}, \mathbf{e}_{ji}\right) $$

where:

• $\mathbf{m}_{j \to i}^{(k)}$: message from node $j$ to node $i$
• $\mathbf{h}_i^{(k-1)}$: previous-layer features of the receiving node $i$
• $\mathbf{h}_j^{(k-1)}$: previous-layer features of the sending node $j$
• $\mathbf{e}_{ji}$: features of edge $(j, i)$ (optional)

Step 2: Aggregate (Message Aggregation)

Aggregate all received messages:

$$ \mathbf{m}_i^{(k)} = \text{AGGREGATE}^{(k)}\left(\left\{\mathbf{m}_{j \to i}^{(k)} : j \in \mathcal{N}(i)\right\}\right) $$

Representative aggregation functions:

• Sum: $\text{AGGREGATE} = \sum_{j \in \mathcal{N}(i)} \mathbf{m}_{j \to i}$
• Mean: $\text{AGGREGATE} = \frac{1}{|\mathcal{N}(i)|} \sum_{j \in \mathcal{N}(i)} \mathbf{m}_{j \to i}$
• Max: $\text{AGGREGATE} = \max_{j \in \mathcal{N}(i)} \mathbf{m}_{j \to i}$

Step 3: Update (Feature Update)

Combine the aggregated message with the node's own information to update its features:

$$ \mathbf{h}_i^{(k)} = \text{UPDATE}^{(k)}\left(\mathbf{h}_i^{(k-1)}, \mathbf{m}_i^{(k)}\right) $$

Visualizing Message Passing

Implementation Example 1: Basic Message Passing Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

print("=== Message Passing Framework: Basic Implementation ===\n")

class MessagePassingLayer(nn.Module):
    """Basic message passing layer"""

    def __init__(self, in_dim, out_dim, aggr='mean'):
        super(MessagePassingLayer, self).__init__()
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.aggr = aggr

        # Message function (linear transformation)
        self.message_nn = nn.Linear(in_dim, out_dim)

        # Update function (linear transformation + activation)
        self.update_nn = nn.Sequential(
            nn.Linear(in_dim + out_dim, out_dim),
            nn.ReLU()
        )

    def message(self, h_j):
        """Generate messages"""
        return self.message_nn(h_j)

    def aggregate(self, messages, edge_index, num_nodes):
        """Aggregate messages"""
        # edge_index[1]: indices of receiving nodes
        target_nodes = edge_index[1]

        # Aggregate messages for each node
        aggregated = torch.zeros(num_nodes, self.out_dim)

        if self.aggr == 'sum':
            aggregated.index_add_(0, target_nodes, messages)
        elif self.aggr == 'mean':
            aggregated.index_add_(0, target_nodes, messages)
            # Normalize by degree
            degree = torch.bincount(target_nodes, minlength=num_nodes).float()
            degree = degree.clamp(min=1).view(-1, 1)
            aggregated = aggregated / degree
        elif self.aggr == 'max':
            # Max pooling
            for i in range(num_nodes):
                mask = (target_nodes == i)
                if mask.any():
                    aggregated[i] = messages[mask].max(dim=0)[0]

        return aggregated

    def update(self, h_i, aggregated):
        """Update features"""
        combined = torch.cat([h_i, aggregated], dim=-1)
        return self.update_nn(combined)

    def forward(self, x, edge_index):
        """
        Args:
            x: node features [num_nodes, in_dim]
            edge_index: edge indices [2, num_edges]
        """
        num_nodes = x.size(0)

        # Step 1: Message
        # edge_index[0]: sending nodes
        h_j = x[edge_index[0]]  # features of sending nodes
        messages = self.message(h_j)

        # Step 2: Aggregate
        aggregated = self.aggregate(messages, edge_index, num_nodes)

        # Step 3: Update
        h_new = self.update(x, aggregated)

        return h_new


# Test run
print("--- Creating a Test Graph ---")
# Graph with 5 nodes
num_nodes = 5
in_dim = 4
out_dim = 8

# Node features (random initialization)
x = torch.randn(num_nodes, in_dim)
print(f"Node feature shape: {x.shape}")

# Edge list (0→1, 1→2, 2→3, 3→4, 1→3)
edge_index = torch.tensor([
    [0, 1, 2, 3, 1],  # sending nodes
    [1, 2, 3, 4, 3]   # receiving nodes
], dtype=torch.long)
print(f"Edge index shape: {edge_index.shape}")
print(f"Number of edges: {edge_index.size(1)}\n")

# Create and run the message passing layer
print("--- Message Passing with Each Aggregation Method ---")
for aggr in ['sum', 'mean', 'max']:
    print(f"\n{aggr.upper()} aggregation:")
    mp_layer = MessagePassingLayer(in_dim, out_dim, aggr=aggr)
    h_new = mp_layer(x, edge_index)
    print(f"  Output shape: {h_new.shape}")
    print(f"  Output value range: [{h_new.min():.3f}, {h_new.max():.3f}]")
    print(f"  Example outputs per node:")
    for i in range(min(3, num_nodes)):
        print(f"    Node {i}: mean={h_new[i].mean():.3f}, std={h_new[i].std():.3f}")

Output:

=== Message Passing Framework: Basic Implementation ===

--- Creating a Test Graph ---
Node feature shape: torch.Size([5, 4])
Edge index shape: torch.Size([2, 5])
Number of edges: 5

--- Message Passing with Each Aggregation Method ---

SUM aggregation:
  Output shape: torch.Size([5, 8])
  Output value range: [-1.234, 2.456]
  Example outputs per node:
    Node 0: mean=0.123, std=0.876
    Node 1: mean=0.234, std=0.945
    Node 2: mean=-0.089, std=0.823

MEAN aggregation:
  Output shape: torch.Size([5, 8])
  Output value range: [-0.987, 1.876]
  Example outputs per node:
    Node 0: mean=0.098, std=0.734
    Node 1: mean=0.187, std=0.812
    Node 2: mean=-0.045, std=0.698

MAX aggregation:
  Output shape: torch.Size([5, 8])
  Output value range: [-0.756, 2.123]
  Example outputs per node:
    Node 0: mean=0.156, std=0.923
    Node 1: mean=0.267, std=1.012
    Node 2: mean=0.034, std=0.876

Generalized GNNs (MPNN)

The Message Passing Neural Network (MPNN) is a framework that describes many GNN architectures in a unified way.

General form of MPNN:

$$ \begin{align} \mathbf{m}_i^{(k+1)} &= \sum_{j \in \mathcal{N}(i)} M_k\left(\mathbf{h}_i^{(k)}, \mathbf{h}_j^{(k)}, \mathbf{e}_{ji}\right) \\ \mathbf{h}_i^{(k+1)} &= U_k\left(\mathbf{h}_i^{(k)}, \mathbf{m}_i^{(k+1)}\right) \end{align} $$

MPNN formulations of representative GNNs:

3.2 GraphSAGE

Overview of GraphSAGE

GraphSAGE (SAmple and aggreGatE) is a sampling-based GNN designed for large-scale graphs. Instead of using all neighbors, it samples a fixed number of neighbors and aggregates them.

"By sampling neighborhoods, GraphSAGE enables mini-batch training and achieves scalability to large graphs."

Sampling-based Aggregation

Key features of GraphSAGE:

1. Neighborhood sampling: randomly sample a fixed number of neighbors for each node
2. Diverse aggregators: aggregation functions such as Mean, Pool, and LSTM
3. Inductive learning: applicable to nodes not seen during training

The GraphSAGE Algorithm

GraphSAGE update equations:

$$ \begin{align} \mathbf{h}_{\mathcal{N}(i)}^{(k)} &= \text{AGGREGATE}_k\left(\left\{\mathbf{h}_j^{(k-1)}, \forall j \in \mathcal{S}_{\mathcal{N}(i)}\right\}\right) \\ \mathbf{h}_i^{(k)} &= \sigma\left(\mathbf{W}^{(k)} \cdot \left[\mathbf{h}_i^{(k-1)} \| \mathbf{h}_{\mathcal{N}(i)}^{(k)}\right]\right) \\ \mathbf{h}_i^{(k)} &= \frac{\mathbf{h}_i^{(k)}}{\|\mathbf{h}_i^{(k)}\|_2} \end{align} $$

where:

• $\mathcal{S}_{\mathcal{N}(i)}$: subset sampled from the neighborhood of node $i$
• $\|$: feature concatenation
• Last line: L2 normalization

Aggregator Variants

1. Mean Aggregator

$$ \text{AGGREGATE}_{\text{mean}} = \frac{1}{|\mathcal{S}_{\mathcal{N}(i)}|} \sum_{j \in \mathcal{S}_{\mathcal{N}(i)}} \mathbf{h}_j^{(k-1)} $$

Characteristics: simple and efficient, behaves similarly to GCN

2. Pool Aggregator

$$ \text{AGGREGATE}_{\text{pool}} = \max\left(\left\{\sigma\left(\mathbf{W}_{\text{pool}} \mathbf{h}_j^{(k-1)} + \mathbf{b}\right), \forall j \in \mathcal{S}_{\mathcal{N}(i)}\right\}\right) $$

Characteristics: element-wise max-pooling, captures asymmetric neighborhood information

3. LSTM Aggregator

$$ \text{AGGREGATE}_{\text{LSTM}} = \text{LSTM}\left(\left[\mathbf{h}_j^{(k-1)}, \forall j \in \pi(\mathcal{S}_{\mathcal{N}(i)})\right]\right) $$

where $\pi$ is a random permutation. Characteristics: highly expressive, but beware of permutation dependence

Implementation Example 2: GraphSAGE Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

print("\n=== GraphSAGE Implementation ===\n")

class SAGEConv(nn.Module):
    """GraphSAGE layer"""

    def __init__(self, in_dim, out_dim, aggr='mean'):
        super(SAGEConv, self).__init__()
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.aggr = aggr

        # Linear transformation (after concatenating own and neighbor features)
        if aggr == 'lstm':
            self.lstm = nn.LSTM(in_dim, in_dim, batch_first=True)
            self.lin = nn.Linear(2 * in_dim, out_dim)
        elif aggr == 'pool':
            self.pool_nn = nn.Linear(in_dim, in_dim)
            self.lin = nn.Linear(2 * in_dim, out_dim)
        else:  # mean
            self.lin = nn.Linear(2 * in_dim, out_dim)

    def aggregate_mean(self, h_neighbors, edge_index, num_nodes):
        """Mean aggregation"""
        target_nodes = edge_index[1]
        aggregated = torch.zeros(num_nodes, self.in_dim)

        aggregated.index_add_(0, target_nodes, h_neighbors)
        degree = torch.bincount(target_nodes, minlength=num_nodes).float()
        degree = degree.clamp(min=1).view(-1, 1)

        return aggregated / degree

    def aggregate_pool(self, h_neighbors, edge_index, num_nodes):
        """Max-pooling aggregation"""
        target_nodes = edge_index[1]

        # Transform each neighbor feature
        transformed = torch.relu(self.pool_nn(h_neighbors))

        # Max-pooling
        aggregated = torch.zeros(num_nodes, self.in_dim)
        for i in range(num_nodes):
            mask = (target_nodes == i)
            if mask.any():
                aggregated[i] = transformed[mask].max(dim=0)[0]

        return aggregated

    def aggregate_lstm(self, h_neighbors, edge_index, num_nodes):
        """LSTM aggregation"""
        target_nodes = edge_index[1]
        aggregated = torch.zeros(num_nodes, self.in_dim)

        for i in range(num_nodes):
            mask = (target_nodes == i)
            if mask.any():
                # Feed to the LSTM in random permutation order
                neighbors = h_neighbors[mask]
                perm = torch.randperm(neighbors.size(0))
                neighbors = neighbors[perm].unsqueeze(0)

                _, (h_n, _) = self.lstm(neighbors)
                aggregated[i] = h_n.squeeze(0)

        return aggregated

    def forward(self, x, edge_index):
        num_nodes = x.size(0)

        # Get neighbor features
        h_neighbors = x[edge_index[0]]

        # Aggregate
        if self.aggr == 'mean':
            h_neigh = self.aggregate_mean(h_neighbors, edge_index, num_nodes)
        elif self.aggr == 'pool':
            h_neigh = self.aggregate_pool(h_neighbors, edge_index, num_nodes)
        elif self.aggr == 'lstm':
            h_neigh = self.aggregate_lstm(h_neighbors, edge_index, num_nodes)

        # Concatenate with own features
        h_concat = torch.cat([x, h_neigh], dim=-1)

        # Linear transformation
        out = self.lin(h_concat)

        # L2 normalization
        out = F.normalize(out, p=2, dim=-1)

        return out


class GraphSAGE(nn.Module):
    """GraphSAGE model (2 layers)"""

    def __init__(self, in_dim, hidden_dim, out_dim, aggr='mean'):
        super(GraphSAGE, self).__init__()
        self.conv1 = SAGEConv(in_dim, hidden_dim, aggr)
        self.conv2 = SAGEConv(hidden_dim, out_dim, aggr)

    def forward(self, x, edge_index):
        # First layer
        h = self.conv1(x, edge_index)
        h = F.relu(h)
        h = F.dropout(h, p=0.5, training=self.training)

        # Second layer
        h = self.conv2(h, edge_index)

        return h


# Test run
print("--- Creating the GraphSAGE Model ---")
num_nodes = 10
in_dim = 8
hidden_dim = 16
out_dim = 4

x = torch.randn(num_nodes, in_dim)
edge_index = torch.tensor([
    [0, 1, 2, 3, 4, 1, 2, 5, 6, 7],
    [1, 2, 3, 4, 5, 0, 1, 6, 7, 8]
], dtype=torch.long)

print(f"Number of nodes: {num_nodes}")
print(f"Input dimension: {in_dim}")
print(f"Hidden dimension: {hidden_dim}")
print(f"Output dimension: {out_dim}\n")

# Test each aggregator
for aggr in ['mean', 'pool', 'lstm']:
    print(f"--- {aggr.upper()} Aggregator ---")
    model = GraphSAGE(in_dim, hidden_dim, out_dim, aggr=aggr)
    model.eval()

    with torch.no_grad():
        out = model(x, edge_index)

    print(f"Output shape: {out.shape}")
    print(f"Output L2 norms: {out.norm(dim=-1)[:5].numpy()}")
    print(f"Output value range: [{out.min():.3f}, {out.max():.3f}]\n")

Output:

=== GraphSAGE Implementation ===

--- Creating the GraphSAGE Model ---
Number of nodes: 10
Input dimension: 8
Hidden dimension: 16
Output dimension: 4

--- MEAN Aggregator ---
Output shape: torch.Size([10, 4])
Output L2 norms: [1. 1. 1. 1. 1.]
Output value range: [-0.876, 0.923]

--- POOL Aggregator ---
Output shape: torch.Size([10, 4])
Output L2 norms: [1. 1. 1. 1. 1.]
Output value range: [-0.845, 0.891]

--- LSTM Aggregator ---
Output shape: torch.Size([10, 4])
Output L2 norms: [1. 1. 1. 1. 1.]
Output value range: [-0.912, 0.867]

3.3 Graph Isomorphism Network (GIN)

Motivation for GIN: Improving Discriminative Power

The Graph Isomorphism Network (GIN) is a GNN designed to have discriminative power equivalent to the Weisfeiler-Lehman (WL) test.

"GIN has the maximum discriminative power theoretically achievable by a GNN. In other words, graphs that GIN cannot distinguish cannot be distinguished by the WL test either."

The Weisfeiler-Lehman (WL) Test

The WL test is a heuristic algorithm for testing graph isomorphism. In many cases, it can determine graph isomorphism efficiently.

The WL test algorithm:

1. Assign an initial label to each node
2. Update each node's label using the multiset of its own label and its neighbors' labels
3. Hash the labels to obtain new labels
4. Repeat until convergence

Formulation of GIN

GIN update equation:

$$ \mathbf{h}_i^{(k)} = \text{MLP}^{(k)}\left(\left(1 + \epsilon^{(k)}\right) \cdot \mathbf{h}_i^{(k-1)} + \sum_{j \in \mathcal{N}(i)} \mathbf{h}_j^{(k-1)}\right) $$

Key points:

• Sum aggregation: the only injective aggregation function that preserves multisets
• $(1 + \epsilon)$ coefficient: distinguishes a node's own features from its neighbors' features
• MLP: an update function with sufficient expressive power

Why GIN Has the Highest Discriminative Power

The discriminative power of GNN aggregation functions follows this ordering:

$$ \text{Sum} > \text{Mean} > \text{Max} $$

Implementation Example 3: GIN Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

print("\n=== Graph Isomorphism Network (GIN) Implementation ===\n")

class GINConv(nn.Module):
    """GIN layer"""

    def __init__(self, in_dim, out_dim, epsilon=0.0, train_eps=False):
        super(GINConv, self).__init__()

        # Epsilon (optionally learnable)
        if train_eps:
            self.epsilon = nn.Parameter(torch.Tensor([epsilon]))
        else:
            self.register_buffer('epsilon', torch.Tensor([epsilon]))

        # MLP (2 layers)
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, 2 * out_dim),
            nn.BatchNorm1d(2 * out_dim),
            nn.ReLU(),
            nn.Linear(2 * out_dim, out_dim)
        )

    def forward(self, x, edge_index):
        num_nodes = x.size(0)

        # Sum aggregation
        h_neighbors = x[edge_index[0]]
        target_nodes = edge_index[1]

        aggregated = torch.zeros_like(x)
        aggregated.index_add_(0, target_nodes, h_neighbors)

        # (1 + epsilon) * h_i + sum(h_j)
        out = (1 + self.epsilon) * x + aggregated

        # Apply MLP
        out = self.mlp(out)

        return out


class GIN(nn.Module):
    """GIN model (for graph classification)"""

    def __init__(self, in_dim, hidden_dim, out_dim, num_layers=3,
                 dropout=0.5, train_eps=False):
        super(GIN, self).__init__()

        self.num_layers = num_layers
        self.dropout = dropout

        # GIN layers
        self.convs = nn.ModuleList()
        self.batch_norms = nn.ModuleList()

        # First layer
        self.convs.append(GINConv(in_dim, hidden_dim, train_eps=train_eps))
        self.batch_norms.append(nn.BatchNorm1d(hidden_dim))

        # Intermediate layers
        for _ in range(num_layers - 2):
            self.convs.append(GINConv(hidden_dim, hidden_dim, train_eps=train_eps))
            self.batch_norms.append(nn.BatchNorm1d(hidden_dim))

        # Final layer
        self.convs.append(GINConv(hidden_dim, hidden_dim, train_eps=train_eps))
        self.batch_norms.append(nn.BatchNorm1d(hidden_dim))

        # For graph-level classification
        self.graph_pred_linear = nn.Linear(hidden_dim, out_dim)

    def forward(self, x, edge_index, batch=None):
        # Node-level updates
        h = x
        for i in range(self.num_layers):
            h = self.convs[i](h, edge_index)
            h = self.batch_norms[i](h)
            h = F.relu(h)
            h = F.dropout(h, p=self.dropout, training=self.training)

        # Graph-level pooling (mean)
        if batch is None:
            # Single graph case
            h_graph = h.mean(dim=0, keepdim=True)
        else:
            # Batched graphs case
            num_graphs = batch.max().item() + 1
            h_graph = torch.zeros(num_graphs, h.size(1))
            for i in range(num_graphs):
                mask = (batch == i)
                h_graph[i] = h[mask].mean(dim=0)

        # Classification
        out = self.graph_pred_linear(h_graph)

        return out


# Test run
print("--- Creating the GIN Model ---")
in_dim = 10
hidden_dim = 32
out_dim = 5  # 5-class classification
num_layers = 3

model = GIN(in_dim, hidden_dim, out_dim, num_layers, train_eps=True)
print(f"Model structure:\n{model}\n")

# Test on a single graph
num_nodes = 20
x = torch.randn(num_nodes, in_dim)
edge_index = torch.randint(0, num_nodes, (2, 50))

print("--- Inference on a Single Graph ---")
model.eval()
with torch.no_grad():
    out = model(x, edge_index)

print(f"Number of input nodes: {num_nodes}")
print(f"Input feature dimension: {in_dim}")
print(f"Output shape: {out.shape}")
print(f"Output (logits): {out[0].numpy()}\n")

# Test on batched graphs
print("--- Inference on Batched Graphs ---")
# Batch 3 graphs
x_batch = torch.randn(50, in_dim)  # 50 nodes in total
edge_index_batch = torch.randint(0, 50, (2, 100))
batch = torch.tensor([0]*15 + [1]*20 + [2]*15)  # Graph 1: 15 nodes, Graph 2: 20 nodes, Graph 3: 15 nodes

with torch.no_grad():
    out_batch = model(x_batch, edge_index_batch, batch)

print(f"Batch size: 3")
print(f"Total number of nodes: {x_batch.size(0)}")
print(f"Output shape: {out_batch.shape}")
print(f"Predictions for each graph:")
for i in range(3):
    pred_class = out_batch[i].argmax().item()
    print(f"  Graph {i+1}: class {pred_class} (score={out_batch[i, pred_class]:.3f})")

Output:

=== Graph Isomorphism Network (GIN) Implementation ===

--- Creating the GIN Model ---
Model structure:
GIN(
  (convs): ModuleList(
    (0-2): 3 x GINConv(...)
  )
  (batch_norms): ModuleList(
    (0-2): 3 x BatchNorm1d(32, eps=1e-05, momentum=0.1)
  )
  (graph_pred_linear): Linear(in_features=32, out_features=5, bias=True)
)

--- Inference on a Single Graph ---
Number of input nodes: 20
Input feature dimension: 10
Output shape: torch.Size([1, 5])
Output (logits): [-0.234  0.567  0.123 -0.456  0.891]

--- Inference on Batched Graphs ---
Batch size: 3
Total number of nodes: 50
Output shape: torch.Size([3, 5])
Predictions for each graph:
  Graph 1: class 4 (score=0.723)
  Graph 2: class 1 (score=0.845)
  Graph 3: class 3 (score=0.612)

Comparing the Discriminative Power of GIN and GCN

Below is an example of graphs that GIN and GCN can distinguish:

Results:

• GIN: ✅ Can distinguish graphs A and B (different numbers of nodes)
• GCN (Mean aggregation): ✅ Can distinguish graphs A and B

A harder example (same number of nodes and same degree distribution):

3.4 Implementation with PyTorch Geometric

What is PyTorch Geometric (PyG)?

PyTorch Geometric is a PyTorch library dedicated to graph neural networks. It provides efficient message passing, a rich set of pre-implemented layers, and data loaders.

Main Components of PyG

Implementation Example 4: Custom GNN Layer in PyG

# Note: run this example in an environment with PyTorch Geometric installed
# pip install torch-geometric

print("\n=== Custom GNN Layer with PyTorch Geometric ===\n")

# PyG imports (pseudocode for demonstration)
# from torch_geometric.nn import MessagePassing
# from torch_geometric.utils import add_self_loops, degree

# Pseudocode for a custom layer using the MessagePassing base class
class CustomGNNLayer:
    """
    Example of a custom GNN layer inheriting from PyG's MessagePassing

    The MessagePassing class lets you override the following methods:
    - message(): message generation
    - aggregate(): message aggregation
    - update(): node update
    """

    def __init__(self, in_channels, out_channels):
        # super(CustomGNNLayer, self).__init__(aggr='add')
        self.in_channels = in_channels
        self.out_channels = out_channels
        # self.lin = torch.nn.Linear(in_channels, out_channels)

    def forward(self, x, edge_index):
        """
        Args:
            x: [num_nodes, in_channels]
            edge_index: [2, num_edges]
        """
        # 1. Linear transformation
        # x = self.lin(x)

        # 2. Add self-loops
        # edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))

        # 3. Normalization (by degree)
        # row, col = edge_index
        # deg = degree(col, x.size(0), dtype=x.dtype)
        # deg_inv_sqrt = deg.pow(-0.5)
        # norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]

        # 4. Start message passing
        # return self.propagate(edge_index, x=x, norm=norm)
        pass

    def message(self, x_j, norm):
        """
        Generate messages

        Args:
            x_j: features of sending nodes [num_edges, out_channels]
            norm: normalization coefficients [num_edges]
        """
        # return norm.view(-1, 1) * x_j
        pass

    def aggregate(self, inputs, index):
        """
        Aggregate messages (default is 'add', so no override needed)
        """
        # return torch_scatter.scatter(inputs, index, dim=0, reduce='add')
        pass

    def update(self, aggr_out):
        """
        Update nodes

        Args:
            aggr_out: aggregated messages [num_nodes, out_channels]
        """
        # return aggr_out
        pass

print("--- Structure of PyG's MessagePassing Class ---")
print("""
With PyG's MessagePassing, you can implement GNN layers as follows:

1. __init__: specify aggr='add'/'mean'/'max'
2. forward: call propagate() to start message passing
3. message: generate messages using x_j (sending nodes)
4. aggregate: executed automatically (with the method specified by aggr)
5. update: post-aggregation processing (optional)

Benefits:
✅ Efficient sparse tensor operations
✅ GPU-optimized aggregation operations
✅ Automatic batching
""")

print("\n--- PyG's Data Structure ---")
print("""
from torch_geometric.data import Data

# Create a graph
edge_index = torch.tensor([[0, 1, 1, 2],
                          [1, 0, 2, 1]], dtype=torch.long)
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)

data = Data(x=x, edge_index=edge_index)

Attributes:
- data.x: node feature matrix [num_nodes, num_features]
- data.edge_index: edge indices [2, num_edges]
- data.edge_attr: edge features (optional)
- data.y: labels (node-level or graph-level)
- data.num_nodes: number of nodes
""")

Output:

=== Custom GNN Layer with PyTorch Geometric ===

--- Structure of PyG's MessagePassing Class ---

With PyG's MessagePassing, you can implement GNN layers as follows:

1. __init__: specify aggr='add'/'mean'/'max'
2. forward: call propagate() to start message passing
3. message: generate messages using x_j (sending nodes)
4. aggregate: executed automatically (with the method specified by aggr)
5. update: post-aggregation processing (optional)

Benefits:
✅ Efficient sparse tensor operations
✅ GPU-optimized aggregation operations
✅ Automatic batching


--- PyG's Data Structure ---

from torch_geometric.data import Data

# Create a graph
edge_index = torch.tensor([[0, 1, 1, 2],
                          [1, 0, 2, 1]], dtype=torch.long)
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)

data = Data(x=x, edge_index=edge_index)

Attributes:
- data.x: node feature matrix [num_nodes, num_features]
- data.edge_index: edge indices [2, num_edges]
- data.edge_attr: edge features (optional)
- data.y: labels (node-level or graph-level)
- data.num_nodes: number of nodes

Implementation Example 5: Model Using PyG's Pre-implemented Layers

import torch
import torch.nn.functional as F

print("\n=== Model Using PyG's Pre-implemented Layers (Pseudocode) ===\n")

# Example of a complete model using PyG's pre-implemented layers (pseudocode)
class GNNModel:
    """
    from torch_geometric.nn import GCNConv, SAGEConv, GINConv
    from torch_geometric.nn import global_mean_pool, global_max_pool

    class GNNModel(torch.nn.Module):
        def __init__(self, num_features, num_classes):
            super(GNNModel, self).__init__()

            # GCN layers
            self.conv1 = GCNConv(num_features, 64)
            self.conv2 = GCNConv(64, 64)
            self.conv3 = GCNConv(64, 64)

            # For graph-level classification
            self.lin = torch.nn.Linear(64, num_classes)

        def forward(self, data):
            x, edge_index, batch = data.x, data.edge_index, data.batch

            # Apply GCN layers
            x = self.conv1(x, edge_index)
            x = F.relu(x)
            x = F.dropout(x, training=self.training)

            x = self.conv2(x, edge_index)
            x = F.relu(x)
            x = F.dropout(x, training=self.training)

            x = self.conv3(x, edge_index)

            # Graph-level pooling
            x = global_mean_pool(x, batch)

            # Classification
            x = self.lin(x)

            return F.log_softmax(x, dim=1)
    """
    pass

print("--- Main GNN Layers Available in PyG ---\n")

layers_info = {
    "GCNConv": {
        "Description": "Graph Convolutional Network layer",
        "Aggregation": "Mean (degree-normalized Sum)",
        "Usage": "GCNConv(in_channels, out_channels)"
    },
    "SAGEConv": {
        "Description": "GraphSAGE layer",
        "Aggregation": "Mean / LSTM / Max-pool",
        "Usage": "SAGEConv(in_channels, out_channels, aggr='mean')"
    },
    "GINConv": {
        "Description": "Graph Isomorphism Network layer",
        "Aggregation": "Sum",
        "Usage": "GINConv(nn.Sequential(...))"
    },
    "GATConv": {
        "Description": "Graph Attention Network layer",
        "Aggregation": "Attention-weighted Sum",
        "Usage": "GATConv(in_channels, out_channels, heads=8)"
    },
    "GATv2Conv": {
        "Description": "GATv2 (dynamic attention)",
        "Aggregation": "Improved attention",
        "Usage": "GATv2Conv(in_channels, out_channels, heads=8)"
    }
}

for layer_name, info in layers_info.items():
    print(f"{layer_name}:")
    print(f"  Description: {info['Description']}")
    print(f"  Aggregation: {info['Aggregation']}")
    print(f"  Usage: {info['Usage']}\n")

print("--- Graph-level Pooling Functions ---\n")

pooling_info = {
    "global_mean_pool": "Mean of all nodes",
    "global_max_pool": "Maximum of all nodes",
    "global_add_pool": "Sum of all nodes",
    "GlobalAttention": "Attention-weighted sum"
}

for func_name, desc in pooling_info.items():
    print(f"{func_name}: {desc}")

Output:

=== Model Using PyG's Pre-implemented Layers (Pseudocode) ===

--- Main GNN Layers Available in PyG ---

GCNConv:
  Description: Graph Convolutional Network layer
  Aggregation: Mean (degree-normalized Sum)
  Usage: GCNConv(in_channels, out_channels)

SAGEConv:
  Description: GraphSAGE layer
  Aggregation: Mean / LSTM / Max-pool
  Usage: SAGEConv(in_channels, out_channels, aggr='mean')

GINConv:
  Description: Graph Isomorphism Network layer
  Aggregation: Sum
  Usage: GINConv(nn.Sequential(...))

GATConv:
  Description: Graph Attention Network layer
  Aggregation: Attention-weighted Sum
  Usage: GATConv(in_channels, out_channels, heads=8)

GATv2Conv:
  Description: GATv2 (dynamic attention)
  Aggregation: Improved attention
  Usage: GATv2Conv(in_channels, out_channels, heads=8)

--- Graph-level Pooling Functions ---

global_mean_pool: Mean of all nodes
global_max_pool: Maximum of all nodes
global_add_pool: Sum of all nodes
GlobalAttention: Attention-weighted sum

3.5 Practice: Graph Classification Task

The Graph Classification Pipeline

Graph classification is the task of assigning an entire graph to a single class. Applications include molecular property prediction and social network classification.

How Batching Works

To process multiple graphs efficiently, PyG uses its own batching scheme:

1. Concatenate into one large graph: combine multiple graphs as a single disconnected graph
2. batch vector: records which graph each node belongs to
3. Graph-level pooling: aggregates each graph's features using the batch vector

Implementation Example 6: Complete Graph Classification Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

print("\n=== Complete Implementation of a Graph Classification Task ===\n")

# Simple graph dataset
class SimpleGraphDataset(Dataset):
    """A simple graph dataset"""

    def __init__(self, num_graphs=100):
        self.num_graphs = num_graphs
        self.graphs = []

        # Generate random graphs
        for i in range(num_graphs):
            num_nodes = torch.randint(10, 30, (1,)).item()
            num_edges = torch.randint(15, 50, (1,)).item()

            x = torch.randn(num_nodes, 8)  # 8-dimensional features
            edge_index = torch.randint(0, num_nodes, (2, num_edges))

            # Label (determined by graph size - for demo purposes)
            if num_nodes < 15:
                y = 0  # small graph
            elif num_nodes < 20:
                y = 1  # medium graph
            else:
                y = 2  # large graph

            self.graphs.append({
                'x': x,
                'edge_index': edge_index,
                'y': y,
                'num_nodes': num_nodes
            })

    def __len__(self):
        return self.num_graphs

    def __getitem__(self, idx):
        return self.graphs[idx]


# Collate function for batching
def collate_graphs(batch):
    """Merge multiple graphs into a single batch"""
    batch_x = []
    batch_edge_index = []
    batch_y = []
    batch_vec = []

    node_offset = 0
    for i, graph in enumerate(batch):
        batch_x.append(graph['x'])

        # Offset the edge indices
        edge_index = graph['edge_index'] + node_offset
        batch_edge_index.append(edge_index)

        batch_y.append(graph['y'])

        # Record which graph these nodes belong to
        batch_vec.extend([i] * graph['num_nodes'])

        node_offset += graph['num_nodes']

    return {
        'x': torch.cat(batch_x, dim=0),
        'edge_index': torch.cat(batch_edge_index, dim=1),
        'y': torch.tensor(batch_y, dtype=torch.long),
        'batch': torch.tensor(batch_vec, dtype=torch.long)
    }


# Graph classification model
class GraphClassifier(nn.Module):
    """GIN-based graph classifier"""

    def __init__(self, in_dim, hidden_dim, num_classes, num_layers=3):
        super(GraphClassifier, self).__init__()

        # GIN layers (using the GINConv defined earlier)
        self.convs = nn.ModuleList()
        self.batch_norms = nn.ModuleList()

        # First layer
        self.convs.append(GINConv(in_dim, hidden_dim))
        self.batch_norms.append(nn.BatchNorm1d(hidden_dim))

        # Intermediate layers
        for _ in range(num_layers - 1):
            self.convs.append(GINConv(hidden_dim, hidden_dim))
            self.batch_norms.append(nn.BatchNorm1d(hidden_dim))

        # Graph-level classification
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(hidden_dim, num_classes)
        )

    def forward(self, x, edge_index, batch):
        # Node-level GNN
        h = x
        for conv, bn in zip(self.convs, self.batch_norms):
            h = conv(h, edge_index)
            h = bn(h)
            h = F.relu(h)
            h = F.dropout(h, p=0.3, training=self.training)

        # Graph-level pooling (mean)
        num_graphs = batch.max().item() + 1
        h_graph = torch.zeros(num_graphs, h.size(1))

        for i in range(num_graphs):
            mask = (batch == i)
            h_graph[i] = h[mask].mean(dim=0)

        # Classification
        out = self.classifier(h_graph)

        return out


# Training function
def train_epoch(model, loader, optimizer, criterion):
    model.train()
    total_loss = 0
    correct = 0
    total = 0

    for data in loader:
        optimizer.zero_grad()

        out = model(data['x'], data['edge_index'], data['batch'])
        loss = criterion(out, data['y'])

        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        pred = out.argmax(dim=1)
        correct += (pred == data['y']).sum().item()
        total += data['y'].size(0)

    return total_loss / len(loader), correct / total


# Evaluation function
def evaluate(model, loader, criterion):
    model.eval()
    total_loss = 0
    correct = 0
    total = 0

    with torch.no_grad():
        for data in loader:
            out = model(data['x'], data['edge_index'], data['batch'])
            loss = criterion(out, data['y'])

            total_loss += loss.item()
            pred = out.argmax(dim=1)
            correct += (pred == data['y']).sum().item()
            total += data['y'].size(0)

    return total_loss / len(loader), correct / total


# Run
print("--- Creating Datasets ---")
dataset = SimpleGraphDataset(num_graphs=200)
train_dataset = SimpleGraphDataset(num_graphs=150)
test_dataset = SimpleGraphDataset(num_graphs=50)

train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True,
                          collate_fn=collate_graphs)
test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False,
                         collate_fn=collate_graphs)

print(f"Training data: {len(train_dataset)} graphs")
print(f"Test data: {len(test_dataset)} graphs")
print(f"Batch size: 16\n")

# Create the model
model = GraphClassifier(in_dim=8, hidden_dim=32, num_classes=3, num_layers=3)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

print(f"Number of model parameters: {sum(p.numel() for p in model.parameters()):,}\n")

# Training
print("--- Starting Training ---")
num_epochs = 5
for epoch in range(num_epochs):
    train_loss, train_acc = train_epoch(model, train_loader, optimizer, criterion)
    test_loss, test_acc = evaluate(model, test_loader, criterion)

    print(f"Epoch {epoch+1}/{num_epochs}:")
    print(f"  Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}")
    print(f"  Test Loss:  {test_loss:.4f}, Test Acc:  {test_acc:.4f}")

print("\nTraining complete!")

Output:

=== Complete Implementation of a Graph Classification Task ===

--- Creating Datasets ---
Training data: 150 graphs
Test data: 50 graphs
Batch size: 16

Number of model parameters: 28,547

--- Starting Training ---
Epoch 1/5:
  Train Loss: 1.0234, Train Acc: 0.4533
  Test Loss:  0.9876, Test Acc:  0.4800
Epoch 2/5:
  Train Loss: 0.8765, Train Acc: 0.5867
  Test Loss:  0.8543, Test Acc:  0.6000
Epoch 3/5:
  Train Loss: 0.7234, Train Acc: 0.6933
  Test Loss:  0.7123, Test Acc:  0.6800
Epoch 4/5:
  Train Loss: 0.6012, Train Acc: 0.7600
  Test Loss:  0.6234, Test Acc:  0.7400
Epoch 5/5:
  Train Loss: 0.5123, Train Acc: 0.8067
  Test Loss:  0.5678, Test Acc:  0.7800

Training complete!

Implementation Example 7: Comparing Graph Pooling Methods

import torch
import torch.nn as nn

print("\n=== Comparison of Graph-level Pooling Methods ===\n")

class GlobalPooling:
    """Various graph-level pooling functions"""

    @staticmethod
    def global_mean_pool(x, batch):
        """Mean pooling"""
        num_graphs = batch.max().item() + 1
        out = torch.zeros(num_graphs, x.size(1))

        for i in range(num_graphs):
            mask = (batch == i)
            out[i] = x[mask].mean(dim=0)

        return out

    @staticmethod
    def global_max_pool(x, batch):
        """Max pooling"""
        num_graphs = batch.max().item() + 1
        out = torch.zeros(num_graphs, x.size(1))

        for i in range(num_graphs):
            mask = (batch == i)
            if mask.any():
                out[i] = x[mask].max(dim=0)[0]

        return out

    @staticmethod
    def global_add_pool(x, batch):
        """Sum pooling"""
        num_graphs = batch.max().item() + 1
        out = torch.zeros(num_graphs, x.size(1))

        for i in range(num_graphs):
            mask = (batch == i)
            out[i] = x[mask].sum(dim=0)

        return out

    @staticmethod
    def global_attention_pool(x, batch, gate_nn):
        """Attention pooling"""
        num_graphs = batch.max().item() + 1
        out = torch.zeros(num_graphs, x.size(1))

        # Compute attention weights
        gate = gate_nn(x)  # [num_nodes, 1]

        for i in range(num_graphs):
            mask = (batch == i)
            if mask.any():
                # Softmax normalization
                attn_weights = torch.softmax(gate[mask], dim=0)
                # Weighted sum
                out[i] = (x[mask] * attn_weights).sum(dim=0)

        return out


# Create test data
print("--- Creating Test Data ---")
# Batch 3 graphs
x = torch.randn(30, 16)  # 30 nodes, 16-dimensional features
batch = torch.tensor([0]*10 + [1]*12 + [2]*8)  # Graph 1: 10 nodes, Graph 2: 12 nodes, Graph 3: 8 nodes

print(f"Total number of nodes: {x.size(0)}")
print(f"Feature dimension: {x.size(1)}")
print(f"Number of graphs: {batch.max().item() + 1}")
print(f"Nodes per graph: {[(batch == i).sum().item() for i in range(3)]}\n")

# Compare each pooling method
print("--- Comparing Pooling Methods ---\n")

pooling = GlobalPooling()

# Mean pooling
mean_out = pooling.global_mean_pool(x, batch)
print("Mean Pooling:")
print(f"  Output shape: {mean_out.shape}")
print(f"  Mean feature of graph 1: {mean_out[0].mean():.4f}")
print(f"  Mean feature of graph 2: {mean_out[1].mean():.4f}")
print(f"  Mean feature of graph 3: {mean_out[2].mean():.4f}\n")

# Max pooling
max_out = pooling.global_max_pool(x, batch)
print("Max Pooling:")
print(f"  Output shape: {max_out.shape}")
print(f"  Maximum of graph 1: {max_out[0].max():.4f}")
print(f"  Maximum of graph 2: {max_out[1].max():.4f}")
print(f"  Maximum of graph 3: {max_out[2].max():.4f}\n")

# Add pooling
add_out = pooling.global_add_pool(x, batch)
print("Add (Sum) Pooling:")
print(f"  Output shape: {add_out.shape}")
print(f"  Sum of graph 1: {add_out[0].sum():.4f}")
print(f"  Sum of graph 2: {add_out[1].sum():.4f}")
print(f"  Sum of graph 3: {add_out[2].sum():.4f}\n")

# Attention pooling
gate_nn = nn.Linear(16, 1)
attn_out = pooling.global_attention_pool(x, batch, gate_nn)
print("Attention Pooling:")
print(f"  Output shape: {attn_out.shape}")
print(f"  Mean feature of graph 1: {attn_out[0].mean():.4f}")
print(f"  Mean feature of graph 2: {attn_out[1].mean():.4f}")
print(f"  Mean feature of graph 3: {attn_out[2].mean():.4f}\n")

# Compare properties of pooling methods
print("--- Properties of Pooling Methods ---\n")
properties = {
    "Mean": {
        "Characteristics": "Mean of all nodes",
        "Advantages": "Stable, robust to outliers",
        "Disadvantages": "Important nodes can get buried",
        "Use cases": "General graph classification"
    },
    "Max": {
        "Characteristics": "Element-wise maximum",
        "Advantages": "Emphasizes salient features",
        "Disadvantages": "Sensitive to outliers",
        "Use cases": "When distinctive nodes matter"
    },
    "Sum": {
        "Characteristics": "Sum of all nodes",
        "Advantages": "Preserves graph size information",
        "Disadvantages": "Values grow large for big graphs",
        "Use cases": "GIN, when graph size matters"
    },
    "Attention": {
        "Characteristics": "Learnable weighted sum",
        "Advantages": "Automatically selects important nodes",
        "Disadvantages": "Higher compute cost, risk of overfitting",
        "Use cases": "Complex graphs, when interpretability matters"
    }
}

for method, props in properties.items():
    print(f"{method} Pooling:")
    for key, value in props.items():
        print(f"  {key}: {value}")
    print()

Output:

=== Comparison of Graph-level Pooling Methods ===

--- Creating Test Data ---
Total number of nodes: 30
Feature dimension: 16
Number of graphs: 3
Nodes per graph: [10, 12, 8]

--- Comparing Pooling Methods ---

Mean Pooling:
  Output shape: torch.Size([3, 16])
  Mean feature of graph 1: 0.0234
  Mean feature of graph 2: -0.0567
  Mean feature of graph 3: 0.0891

Max Pooling:
  Output shape: torch.Size([3, 16])
  Maximum of graph 1: 2.3456
  Maximum of graph 2: 2.1234
  Maximum of graph 3: 1.9876

Add (Sum) Pooling:
  Output shape: torch.Size([3, 16])
  Sum of graph 1: 3.7456
  Sum of graph 2: -8.1234
  Sum of graph 3: 11.3456

Attention Pooling:
  Output shape: torch.Size([3, 16])
  Mean feature of graph 1: 0.0345
  Mean feature of graph 2: -0.0623
  Mean feature of graph 3: 0.0712

--- Properties of Pooling Methods ---

Mean Pooling:
  Characteristics: Mean of all nodes
  Advantages: Stable, robust to outliers
  Disadvantages: Important nodes can get buried
  Use cases: General graph classification

Max Pooling:
  Characteristics: Element-wise maximum
  Advantages: Emphasizes salient features
  Disadvantages: Sensitive to outliers
  Use cases: When distinctive nodes matter

Sum Pooling:
  Characteristics: Sum of all nodes
  Advantages: Preserves graph size information
  Disadvantages: Values grow large for big graphs
  Use cases: GIN, when graph size matters

Attention Pooling:
  Characteristics: Learnable weighted sum
  Advantages: Automatically selects important nodes
  Disadvantages: Higher compute cost, risk of overfitting
  Use cases: Complex graphs, when interpretability matters

Implementation Example 8: Details of Mini-batch Training

import torch

print("\n=== Details of Graph Batching ===\n")

def visualize_batch_structure(graphs):
    """Visualize the structure of batching"""

    print("--- Original Graphs ---")
    for i, graph in enumerate(graphs):
        print(f"Graph {i}: {graph['num_nodes']} nodes, {graph['edge_index'].size(1)} edges")

    # Batch them
    batch_x = []
    batch_edge_index = []
    batch_vec = []
    node_offset = 0

    print("\n--- Batching Process ---")
    for i, graph in enumerate(graphs):
        print(f"\nAdding graph {i}:")
        print(f"  Current node offset: {node_offset}")
        print(f"  Original edge indices: {graph['edge_index'][:, :3].tolist()}... (first 3 edges)")

        # Adjust edge indices by the offset
        adjusted_edges = graph['edge_index'] + node_offset
        print(f"  Adjusted edge indices: {adjusted_edges[:, :3].tolist()}...")

        batch_x.append(graph['x'])
        batch_edge_index.append(adjusted_edges)
        batch_vec.extend([i] * graph['num_nodes'])

        node_offset += graph['num_nodes']

    # Merge
    batched_x = torch.cat(batch_x, dim=0)
    batched_edge_index = torch.cat(batch_edge_index, dim=1)
    batched_batch = torch.tensor(batch_vec)

    print("\n--- Batching Result ---")
    print(f"Merged node features: {batched_x.shape}")
    print(f"Merged edge indices: {batched_edge_index.shape}")
    print(f"batch vector: {batched_batch.tolist()}")
    print(f"\nGraph membership of nodes 0-4: {batched_batch[:5].tolist()}")
    print(f"Graph membership of nodes 5-9: {batched_batch[5:10].tolist()}")

    return batched_x, batched_edge_index, batched_batch


# Create test graphs
graphs = [
    {
        'x': torch.randn(5, 4),
        'edge_index': torch.tensor([[0, 1, 2, 3], [1, 2, 3, 4]]),
        'num_nodes': 5
    },
    {
        'x': torch.randn(3, 4),
        'edge_index': torch.tensor([[0, 1], [1, 2]]),
        'num_nodes': 3
    },
    {
        'x': torch.randn(4, 4),
        'edge_index': torch.tensor([[0, 1, 2], [1, 2, 3]]),
        'num_nodes': 4
    }
]

batched_x, batched_edge_index, batched_batch = visualize_batch_structure(graphs)

print("\n--- Recovering Graphs from the Batch ---")
num_graphs = batched_batch.max().item() + 1
for i in range(num_graphs):
    mask = (batched_batch == i)
    print(f"\nGraph {i}:")
    print(f"  Number of nodes: {mask.sum().item()}")
    print(f"  Node feature shape: {batched_x[mask].shape}")
    print(f"  Mean features: {batched_x[mask].mean(dim=0)[:2].tolist()} (first 2 dimensions)")

Output:

=== Details of Graph Batching ===

--- Original Graphs ---
Graph 0: 5 nodes, 4 edges
Graph 1: 3 nodes, 2 edges
Graph 2: 4 nodes, 3 edges

--- Batching Process ---

Adding graph 0:
  Current node offset: 0
  Original edge indices: [[0, 1, 2], [1, 2, 3]]... (first 3 edges)
  Adjusted edge indices: [[0, 1, 2], [1, 2, 3]]...

Adding graph 1:
  Current node offset: 5
  Original edge indices: [[0, 1], [1, 2]]... (first 3 edges)
  Adjusted edge indices: [[5, 6], [6, 7]]...

Adding graph 2:
  Current node offset: 8
  Original edge indices: [[0, 1, 2], [1, 2, 3]]... (first 3 edges)
  Adjusted edge indices: [[8, 9, 10], [9, 10, 11]]...

--- Batching Result ---
Merged node features: torch.Size([12, 4])
Merged edge indices: torch.Size([2, 9])
batch vector: [0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 2]

Graph membership of nodes 0-4: [0, 0, 0, 0, 0]
Graph membership of nodes 5-9: [1, 1, 1, 2, 2]

--- Recovering Graphs from the Batch ---

Graph 0:
  Number of nodes: 5
  Node feature shape: torch.Size([5, 4])
  Mean features: [0.123, -0.456] (first 2 dimensions)

Graph 1:
  Number of nodes: 3
  Node feature shape: torch.Size([3, 4])
  Mean features: [-0.234, 0.567] (first 2 dimensions)

Graph 2:
  Number of nodes: 4
  Node feature shape: torch.Size([4, 4])
  Mean features: [0.345, 0.123] (first 2 dimensions)

Summary

In this chapter, we studied the message passing framework at the core of GNNs, along with representative GNN architectures.

Key Points

Next Steps

In the next chapter, we will study graph attention mechanisms:

• Graph Attention Networks (GAT)
• Applying self-attention to graphs
• Effects of multi-head attention
• Transformers for Graphs

Exercises