Graph neural network layers (spec layer)

This file defines a GCN-style graph convolution layer. Neighbor aggregation uses the shared tensor matrix multiplication and its backward rule.

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.

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 α] [Add α] [Mul α] [Zero α] {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 α] [Add α] [Mul α] [Zero α] {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.