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 this file defines only these two:

• 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} [TorchLean.Storage α] [Context α] {n inDim : ℕ} (A : TorchLean.Tensor α [n, n]) (x : TorchLean.Tensor α [n, inDim]) : TorchLean.Tensor α [n, inDim]

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} [TorchLean.Storage α] [Context α] {n inDim : ℕ} (A : TorchLean.Tensor α [n, n]) (x dY : TorchLean.Tensor α [n, inDim]) : TorchLean.Tensor α [n, n] × TorchLean.Tensor α [n, inDim]

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

structure Spec.GCNLayerSpec (n inDim outDim : ℕ) (α : Type) [TorchLean.Storage α] : 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 : TorchLean.Tensor α [n, n]

  The (possibly normalized) adjacency operator acting on the n nodes.

• W : TorchLean.Tensor α [inDim, outDim]

  The weight matrix mapping input features to output features.

• b : TorchLean.Tensor α [outDim]

  The bias vector, broadcast across nodes.

def Spec.gcnLayerSpec {α : Type} [TorchLean.Storage α] [Context α] {n inDim outDim : ℕ} (layer : GCNLayerSpec n inDim outDim α) (x : TorchLean.Tensor α [n, inDim]) : TorchLean.Tensor α [n, outDim]

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).

structure Spec.GCNLayerParameterGradients (n inDim outDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

Parameter gradients for a GCNLayerSpec.

The adjacency gradient sits with the weight and bias gradients because a GCN layer can be trained with learned edge weights, in which case A really is a parameter; when it is a fixed normalization the field is simply ignored.

• adjacencyGradient : TorchLean.Tensor α [n, n]

  Gradient with respect to the adjacency operator A.

• weightGradient : TorchLean.Tensor α [inDim, outDim]

  Gradient with respect to the weight matrix W.

• biasGradient : TorchLean.Tensor α [outDim]

  Gradient with respect to the bias vector b.

structure Spec.GCNLayerGradients (n inDim outDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

Everything a GCN layer's backward pass produces: the parameter gradients plus the gradient travelling on to the node features.

• parameters : GCNLayerParameterGradients n inDim outDim α

  Gradients for the layer parameters.

• inputGradient : TorchLean.Tensor α [n, inDim]

  Gradient with respect to the node feature matrix X.

def Spec.gcnLayerBackwardSpec {α : Type} [TorchLean.Storage α] [Context α] {n inDim outDim : ℕ} (layer : GCNLayerSpec n inDim outDim α) (x : TorchLean.Tensor α [n, inDim]) (gradOutput : TorchLean.Tensor α [n, outDim]) (h_n : n ≠ 0) : GCNLayerGradients n inDim outDim α

Backward/VJP spec for gcnLayerSpec.