Graph neural network layers (spec layer)

We provide a couple of small, standard GNN building blocks that show up in lots of papers and PyTorch GNN libraries:

• a basic "message passing / neighbor aggregation" primitive, and
• a GCN-style graph convolution layer.

Message passing (the common core idea)

Most GNN layers have the same shape of computation:

• aggregate neighbor features using the graph structure, then
• optionally apply a learnable transformation and a nonlinearity.

In this file the aggregation step is written with a matrix A : (n×n):

Agg(A, H) = A · H.

This captures many common conventions:

• if A is the raw adjacency, you are summing neighbors,
• if A is normalized (e.g. D^{-1/2} (A + I) D^{-1/2}), you are doing the "GCN normalization" flavor,
• if A includes edge weights, you are doing a weighted sum.

GCN layer (one very common choice)

We model a GCN-style layer as:

H' = A · H · W + b

where:

• A : (n×n) is an adjacency-like matrix (often normalized, and often with self-loops),
• H : (n×inDim) are node features,
• W : (inDim×outDim) and b : outDim are trainable parameters.

PyTorch mental picture:

• This is the algebraic core of what libraries like PyTorch Geometric call GCNConv once you pick a concrete choice of A (raw adjacency, D^{-1/2} (A + I) D^{-1/2}, etc.) and batch conventions.

Why only these two right now:

• GCN + plain aggregation are enough to cover a lot of examples and give us something we can reason about cleanly.
• We do plan to add other families (GraphSAGE, GAT, generic MPNNs). Those require more choices (per-edge features, masking/batching conventions, and tie-ins to attention-style ops), so we want to introduce them carefully instead of piling on half-finished variants.

def Spec.messagePassingSpec {α : Type} [Context α] {n inDim : ℕ} (A : Tensor α (Shape.dim n (Shape.dim n Shape.scalar))) (x : Tensor α (Shape.dim n (Shape.dim inDim Shape.scalar))) : Tensor α (Shape.dim n (Shape.dim inDim Shape.scalar))

Neighbor aggregation / message passing via a graph matrix: Agg(A, X) = A · X.

This is the reusable "mix neighbors" step. The semantics are entirely determined by A (raw adjacency, normalized adjacency, weighted adjacency, etc.).

def Spec.messagePassingBackwardSpec {α : Type} [Context α] {n inDim : ℕ} (A : Tensor α (Shape.dim n (Shape.dim n Shape.scalar))) (x dY : Tensor α (Shape.dim n (Shape.dim inDim Shape.scalar))) : Tensor α (Shape.dim n (Shape.dim n Shape.scalar)) × Tensor α (Shape.dim n (Shape.dim inDim Shape.scalar))

Backward/VJP for message_passing_spec: returns (dA, dX).

structure Spec.GCNLayerSpec (n inDim outDim : ℕ) (α : Type) : Type

Parameters/data for a single GCN-style layer.

We bundle A with the layer because many code paths treat A as a fixed input per graph, while others treat it as a parameter (e.g. learned normalization). Keeping it in the record makes both uses explicit.

• A : Tensor α (Shape.dim n (Shape.dim n Shape.scalar))

  A.

• W : Tensor α (Shape.dim inDim (Shape.dim outDim Shape.scalar))

  W.

• b : Tensor α (Shape.dim outDim Shape.scalar)

  b.

def Spec.gcnLayerSpec {α : Type} [Context α] {n inDim outDim : ℕ} (layer : GCNLayerSpec n inDim outDim α) (x : Tensor α (Shape.dim n (Shape.dim inDim Shape.scalar))) : Tensor α (Shape.dim n (Shape.dim outDim Shape.scalar))

Forward spec for a GCN-style layer: Y = A · X · W + b.

Notes:

• The bias b is broadcast across the n nodes (row-wise add).
• Any normalization/self-loop convention belongs in the choice of A supplied to the layer.

Gradients

For the simple GCN-style layer

Y = A · X · W + b

the reverse-mode derivatives are the standard matrix calculus ones:

• dW = (A·X)ᵀ · dY
• db = Σᵢ dYᵢ (sum across the node axis)
• dX = Aᵀ · (dY · Wᵀ)
• dA = (dY · Wᵀ) · Xᵀ

We include dA because in some setups the adjacency/normalization is also:

• treated as an input you want sensitivities for, or
• treated as a parameter (e.g. learned edge weights / learned normalization).

def Spec.gcnLayerBackwardSpec {α : Type} [Context α] {n inDim outDim : ℕ} (layer : GCNLayerSpec n inDim outDim α) (x : Tensor α (Shape.dim n (Shape.dim inDim Shape.scalar))) (grad_output : Tensor α (Shape.dim n (Shape.dim outDim Shape.scalar))) (h_n : n ≠ 0) : Tensor α (Shape.dim n (Shape.dim n Shape.scalar)) × Tensor α (Shape.dim inDim (Shape.dim outDim Shape.scalar)) × Tensor α (Shape.dim outDim Shape.scalar) × Tensor α (Shape.dim n (Shape.dim inDim Shape.scalar))

Backward/VJP spec for gcn_layer_spec.

Returns (dA, dW, db, dX) in that order.