MLP (spec wiring example)

This file defines a small 2-layer MLP by composing SpecChains from module specs:

Linear → ReLU → Linear (optionally followed by a softmax head).

This is intentionally written in a "wiring-first" style: instead of re-implementing matrix multiplications directly, we reuse the spec-layer definitions for Linear and ReLU and compose them through NNModuleSpec / SpecChain. This matches how most PyTorch users think about MLPs: define a few modules, then run a forward pass.

def Examples.mlpSpec {α : Type} [Context α] {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec α inDim hidDim) (l2 : Spec.LinearSpec α hidDim outDim) : ModSpec.SpecChain α (Spec.Shape.dim inDim Spec.Shape.scalar) (Spec.Shape.dim outDim Spec.Shape.scalar)

A 2-layer MLP as a SpecChain:

Linear(inDim → hidDim) then ReLU then Linear(hidDim → outDim).

PyTorch analogy: nn.Sequential(nn.Linear(inDim, hidDim), nn.ReLU(), nn.Linear(hidDim, outDim)).

def Examples.mlpWithSoftmaxSpec {α : Type} [Context α] {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec α inDim hidDim) (l2 : Spec.LinearSpec α hidDim outDim) : ModSpec.SpecChain α (Spec.Shape.dim inDim Spec.Shape.scalar) (Spec.Shape.dim outDim Spec.Shape.scalar)

MLP with a softmax head (Linear → ReLU → Linear → Softmax).

PyTorch analogy: nn.Sequential(..., nn.Softmax(dim=-1)).

Note: this is a shape-safe softmax spec (applied along the last dimension). In PyTorch you choose dim at runtime; here the shape index already tells us what "the last dim" is.

def Examples.mlpForward {α : Type} [Context α] {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec α inDim hidDim) (l2 : Spec.LinearSpec α hidDim outDim) (x : Spec.Tensor α (Spec.Shape.dim inDim Spec.Shape.scalar)) : Spec.Tensor α (Spec.Shape.dim outDim Spec.Shape.scalar)

Run the MLP forward on a single input vector.

def Examples.mlpBackward {α : Type} [Context α] {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec α inDim hidDim) (l2 : Spec.LinearSpec α hidDim outDim) (x : Spec.Tensor α (Spec.Shape.dim inDim Spec.Shape.scalar)) (dLdy : Spec.Tensor α (Spec.Shape.dim outDim Spec.Shape.scalar)) : Spec.Tensor α (Spec.Shape.dim hidDim (Spec.Shape.dim inDim Spec.Shape.scalar)) × Spec.Tensor α (Spec.Shape.dim hidDim Spec.Shape.scalar) × Spec.Tensor α (Spec.Shape.dim outDim (Spec.Shape.dim hidDim Spec.Shape.scalar)) × Spec.Tensor α (Spec.Shape.dim outDim Spec.Shape.scalar) × Spec.Tensor α (Spec.Shape.dim inDim Spec.Shape.scalar)

Backward pass for the 2-layer MLP. Returns (∂L/∂W1, ∂L/∂b1, ∂L/∂W2, ∂L/∂b2, ∂L/∂x).

theorem Examples.mlp_spec_forward_eq {α : Type} [Context α] {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec α inDim hidDim) (l2 : Spec.LinearSpec α hidDim outDim) (x : Spec.Tensor α (Spec.Shape.dim inDim Spec.Shape.scalar)) : (mlpSpec l1 l2).forward x = have z1 := Spec.linearSpec l1 x; have a1 := Activation.reluSpec z1; Spec.linearSpec l2 a1

Sanity lemma: the composed SpecChain forward equals the hand-written Linear → ReLU → Linear computation.

def Examples.mlpOpspec {α : Type} [Context α] {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec α inDim hidDim) (l2 : Spec.LinearSpec α hidDim outDim) : Spec.OpSpec α (Spec.Shape.dim inDim Spec.Shape.scalar) (Spec.Shape.dim outDim Spec.Shape.scalar)

OpSpec for the same 2-layer MLP.

This packaging is convenient for symbolic gradient checks: OpSpec pairs a forward definition with an explicit reverse-mode definition, and it composes cleanly.

def Examples.mlpOpspecBackward {α : Type} [Context α] {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec α inDim hidDim) (l2 : Spec.LinearSpec α hidDim outDim) (x : Spec.Tensor α (Spec.Shape.dim inDim Spec.Shape.scalar)) (dLdy : Spec.Tensor α (Spec.Shape.dim outDim Spec.Shape.scalar)) : Spec.Tensor α (Spec.Shape.dim inDim Spec.Shape.scalar)

Composed backward of the MLP using the OpSpec chain.