Toro

Graph Neural Networks

Toro.GNN provides graph neural network layers following the message-passing framework. Open the namespace to get started:

open Toro
open Toro.NN
open Toro.GNN

Graph Data

Graphs are represented in COO (coordinate) format with a node feature matrix and an edge index tensor:

let x = Tensor.randn ([ 4; 16 ], F32, Cpu)
let edgeIndex = Tensor.ofList ([ [ 0L; 1L; 2L; 3L ]; [ 1L; 2L; 3L; 0L ] ], Cpu)
let g = GraphData.create x edgeIndex

edgeIndex has shape [2, E] where row 0 is the source and row 1 is the target of each edge.

Message Passing

All graph convolution layers follow the message-passing pattern:

  1. Message -- gather features from source nodes along edges.
  2. Aggregate -- scatter messages to target nodes (sum, mean, or max).
  3. Update -- combine aggregated messages with the node's own features.

The MessagePassing module provides the shared aggregate and edgeSoftmax operations.

Convolution Layers

GCNConv

Graph Convolutional Network (Kipf & Welling, 2017). Uses symmetric normalization:

let conv = GCNConv.init 16 32 F32 Cpu
let out = conv.forward (x, edgeIndex)  // [N, 32]

GATConv

Graph Attention Network (Velickovic et al., 2018). Multi-head attention over edges:

let conv = GATConv.init 16 32 4 true 0.2 F32 Cpu
// inChannels=16, outChannels=32, heads=4, concat=true, negativeSlope=0.2
let out = conv.forward (x, edgeIndex)  // [N, 128] (4 heads * 32)

Set concat=false to average over heads instead of concatenating (output: [N, 32]).

SAGEConv

GraphSAGE (Hamilton et al., 2017). Mean-aggregator variant with separate self/neighbor transforms:

let conv = SAGEConv.init 16 32 F32 Cpu
let out = conv.forward (x, edgeIndex)  // [N, 32]

GINConv

Graph Isomorphism Network (Xu et al., 2019). Uses a 2-layer MLP with a learnable epsilon:

let conv = GINConv.init 16 64 32 true F32 Cpu
// inChannels=16, hiddenChannels=64, outChannels=32, trainEps=true
let out = conv.forward (x, edgeIndex)  // [N, 32]

Normalization

GraphNorm

Graph Normalization (Cai et al., 2021). Normalizes node features per graph with a learnable shift:

let norm = GraphNorm.init 32 F32 Cpu
let normalized = norm.forward (out, Some batchVec)

When batch is None, all nodes are treated as belonging to a single graph.

Graph Batching

Combine multiple graphs into a single disconnected graph for batched processing:

let batched = Batch.batch [ g1; g2; g3 ]
// batched.Batch = Some batchVector  ([N_total] mapping nodes to graph index)

Global Pooling

Aggregate node features into graph-level representations:

let batchVec = batched.Batch.Value
let numGraphs = Batch.numGraphs batched
let graphEmb = GlobalPool.globalMeanPool batched.X batchVec numGraphs
// graphEmb: [numGraphs, features]

Available operations: globalMeanPool, globalSumPool, globalMaxPool.

Node Classification Example

open Toro
open Toro.NN
open Toro.GNN

type GcnModel = { Conv1: GCNConv; Conv2: GCNConv }

let forward model (x, edgeIndex) =
    let h = model.Conv1.forward (x, edgeIndex)
    let h = h.relu ()
    model.Conv2.forward (h, edgeIndex)

let conv1 = GCNConv.init inFeatures 16 F32 Cpu
let conv2 = GCNConv.init 16 numClasses F32 Cpu
let model = { Conv1 = conv1; Conv2 = conv2 }

let opt = AdamW.createWithLr 0.01 (Model.trainableVars model)
for _ in 1..200 do
    scoped {
        opt.zeroGrad ()
        let logits = forward model (x, edgeIndex)
        let loss = Loss.crossEntropy logits labels
        loss.backward ()
        opt.step ()
    }