Graph neural networks

Build native GCN, GraphSAGE and GIN models with directed sparse messages, explicit supervision, detached evaluation and retained training snapshots.

API reference

Separate graph structure from supervision

GNN: messages follow directed sparse edges

The source node sends a transformed message to the destination node. Destinations aggregate incoming messages and update their features. Node tasks preserve node rows; graph tasks mean pool rows within each disjoint component before prediction. The representation is an edge list, not a dense nodes by nodes adjacency matrix.

Follow the arrows in the diagram. On a narrow screen, scroll inside the figure or read its caption.

Choose supervision separately from topology

All four nodes remain in the allowed message passing graph. Training indices select target rows 0 and 1; held out evaluation selects target rows 2 and 3. A target selection does not cut an edge. The experiment owner must decide whether the full graph is permitted by a transductive, inductive or temporal split.

Follow the arrows in the diagram. On a narrow screen, scroll inside the figure or read its caption.

Supervised rows and message passing edges are separate. Training indices choose which node or graph targets contribute to the loss. They do not remove neighbors from the graph. A transductive split may retain held out node features and edges; an inductive or temporal study must construct its allowed graph and preprocessing before fitting. Never pass validation/test labels as training features.

This 0.0.7-dev.1 development API belongs to Aner Neural. It uses Aner's Tensor/autodiff engine for trainable GCN, mean GraphSAGE and GIN networks. See shared neural workflows for evaluation, minibatch updates, placement, diagrams and checkpoint ownership. No separate language runtime, downloaded graph library or automatic model download is required.

import aner.nn

let graph = nn.graph(nodes: 4,
    sources: List(0, 1, 2, 3), destinations: List(1, 0, 3, 2))
let features: Tensor = [[-1.0], [-0.8], [0.8], [1.0]]
let labels: Tensor = [[0.0], [0.0], [1.0], [1.0]]
let model = nn.gnn(inputs: 1, hidden: 8, outputs: 2,
    architecture: "sage", layers: 2).task("node_classification").on("cpu")
let fitted = model.fit(graph, features, labels, epochs: 20,
    indices: List(0, 2))
let held_out = fitted.evaluate(graph, features, labels, indices: List(1, 3))
print(held_out.summary())

The example uses a transductive protocol: held out node features and allowed edges remain visible, while only the selected training labels contribute to gradients. It is not an inductive or temporal benchmark. Choose graph boundaries and permitted edges before fitting. A held out link prediction target must be removed from message passing edges; selecting node supervision does not remove edges. Fit feature scaling and vocabulary mappings on the permitted training data, and retain their metadata separately.

Topology and shapes

nn.graph(nodes, sources, destinations, graph_ids: List()) returns immutable GraphData. The two endpoint lists have equal lengths and contain exact Int64 IDs in [0,nodes). Edge sources[e] -> destinations[e] sends a message from the source to the destination. Supply both directions explicitly for an undirected edge. Duplicate edges are separate messages. An empty edge list is valid; isolated nodes remain present. Graph construction admits up to 4M nodes and 16M supplied edges, but operation/tensor budgets can impose much smaller limits.

Optional graph_ids assigns every node to one disjoint graph. IDs are contiguous from zero, each graph contains at least one node, and edges cannot cross graph IDs. This batches differently sized graphs as one sparse disjoint union. graph.nodes(), .edges() and .graph_count() inspect these declarations.

Features are detached [N,F] tensors in stable node ID order. Topology is exact integer metadata, never a floating point tensor or an N × N adjacency matrix. Messages require memory proportional to edge count times hidden width; sparse topology alone does not guarantee a small activation graph. The native nn.segment_sum(values, segment_ids, segments) primitive sums rank two rows and backpropagates by gathering segment gradients, including repeated IDs and empty segments. CPU, Metal and CUDA use their native Tensor gather/scatter operations.

Architectures and tasks

nn.gnn(inputs, hidden, outputs, architecture: "gcn", layers: 2, seed: 42) returns an immutable GNNModel. Feature/output widths are explicit; hidden is the width of each message layer, and a separate linear task head maps to outputs. The default task is node_regression; select classification explicitly. Preparation is lazy until .prepare(), inference or training. .task(...) selects the objective before preparation. .dropout(probability) applies training only inverted dropout after message layers; evaluation disables it.

Architecture	Implemented message/update rule
`gcn`	Add a self loop only where absent; weight each source message by `1/sqrt(in_degree(source)*in_degree(destination))` after these loops; sum, project and ReLU. Existing self loops and duplicate edges retain their multiplicity.
`sage`	Mean incoming neighbors, concatenate with the node's own features, project and ReLU. An empty neighborhood has zero mean; the self projection remains available.
`gin`	Sum incoming neighbors, add `(1+epsilon)` times self features, then a two layer ReLU MLP. Epsilon is a learned scalar per layer.

For directed GCN, the stated incoming degree normalization is deliberate; it does not silently symmetrize a graph. The rules follow the central mechanisms of GCN, GraphSAGE and GIN, with the declared head and activations. They do not claim arbitrary checkpoint compatibility or every paper variant.

Task	Logits/raw values	Targets	Prediction
`node_classification`	`[N,O]`	Exact IDs `[N]`/`[N,1]`, or exact one hot `[N,O]`	`[N,1]` IDs
`node_regression` (default)	`[N,O]`	`[N,O]`	`[N,O]` values
`graph_classification`	`[G,O]`	Exact IDs `[G]`/`[G,1]`, or exact one hot `[G,O]`	`[G,1]` IDs
`graph_regression`	`[G,O]`	`[G,O]`	`[G,O]` values

Graph tasks mean pool node embeddings within each graph, then apply the task head. Classification requires at least two outputs. .logits(graph,x) returns raw scores/values; .probabilities(graph,x) returns classification softmax probabilities; .predict(graph,x) follows the table. All are inference operations returning detached tensors. Regression does not expose probabilities.

Training, evaluation and continuation

model.fit(graph, features, targets, epochs=1, batch=32, rate=0.001,
          seed=42, optimizer="adam", clip_norm=1.0, weight_decay=0.0,
          shuffle=true, indices=List(), execution=...)
model.train_step(graph, features, targets, rate=0.001, optimizer="adam",
                 clip_norm=1.0, weight_decay=0.0, indices=List(), execution=...)
model.evaluate(graph, features, targets, indices=List())

indices selects unique supervised node IDs or graph IDs according to the task; empty means all positions. Targets retain their full position layout. Only selected labels enter the loss. batch counts supervised positions, not nodes removed from message passing. Each batch still computes on the full supplied topology. Cross entropy averages selected class positions; MSE averages selected positions and output features. NeuralEvaluation reports .loss(), .accuracy() (-1 for regression), .samples(), .observations(), .correct(), .objective() and .summary(). Evaluation does not update weights or optimizer state.

Fit supports Adam and SGD, global norm clipping (0 disables it), and decoupled weight decay on projection matrices, excluding biases and GIN epsilon. Rates are finite and positive; decay is nonnegative with rate*weight_decay < 1. Training returns an immutable model; keep the returned value. Matching fits continue saved moments, shuffle epochs and identical graph/data/supervision identity. .reset_optimizer() preserves learned weights and lifetime update counts while clearing moments and data/epoch continuation for a fresh phase.

Use .train_step(...) for an external minibatch loader, a manual learning rate schedule or changing graphs. It performs one global optimizer attempt while retaining compatible moments and cumulative update count. It treats the supplied graph and supervision as an independent batch; it does not enforce whole dataset resume identity or construct a neighbor sampler. Caller owned graph/feature/label alignment remains essential.

.save(path) exclusively creates a versioned, checked checkpoint with topology configuration, learned parameters, optimizer moments, counters and shuffle/data identity. It does not save the input graph, features, labels or preprocessing. nn.load_gnn(path) restores portable CPU storage; use .on(...) deliberately after loading. The format is bounded to 64MiB. .parameter_names(), .parameter(name), .parameter_count(), .updates() and .last_loss() support inspection; returned parameter tensors are detached.

Parallelism and hardware scope

.on("cpu", "Float64"), CPU Float32, native Metal Float32 and CUDA Float32/Float64 use explicit native placement. No activation fallback to CPU is introduced for sparse message passing. Topology metadata and required scalar metrics still travel between host and device. This increment uses strict storage precision; GNN AMP is not configured implicitly.

fit and train_step accept execution: parallel.local(workers: 2) for CPU replicas. Every replica keeps the full graph. Supervised positions are partitioned into unequal/empty shards, gradient/loss sums are reduced, and the coordinator divides once by global supervised observations before one optimizer update. This preserves cross shard neighbors and consistent node addressed dropout. It repeats message passing work and is not graph partitioning.

Optional MPI/NCCL uses the existing CNN launcher/collective contract: one externally launched process per selected visible GPU, identical graph/data, parameter ordering and run fingerprint, coordinated failures, and rank zero checkpoint ownership. .performance(collectives: "mpi") or "nccl" makes that choice explicit. SSH/process parallel.cluster transport remains Dense specific; passing it to these native model families is rejected. Local replica tests and one device tests do not establish physical multi GPU correctness or scaling.

GAT/edge-feature networks, heterogeneous or bipartite graphs, neighbor sampling, partitioned graphs, automatic inductive splits, link prediction heads and model sharding are outside this core implementation. They need their own precise training and topology contracts. Sparse native training, inference, evaluation, checkpoint continuation and supported gradient collectives are implemented here.

Examples and validation

Run node classification or disjoint graph regression. Tests check sparse finite difference gradients, duplicates/isolated nodes, permutation equivariance, disjoint batching, learning, supervision masks, immutable snapshots, checkpoint replay and sample weighted replicas. Device profiles exercise actual forward/backward/optimizer work, with unavailable providers reported as skips.

Aner documentation · Reference and guides