MLP (spec wiring example)

This file defines a 2-layer MLP by composing Spec.Module.Chains from module specs:

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

The file is organized around module wiring rather than re-implementing matrix multiplications directly. Linear and ReLU come from the spec layer and are composed through Spec.Module / Spec.Module.Chain, matching the usual PyTorch workflow: define a few modules, then run a forward pass.

Implementation status

There is no dedicated nn.models builder; the API composes nn.linear and nn.relu directly. The relating theorem is mlp_interp_eq_spec_mlp_forward in NN/GraphSpec/Models/MlpSpecEquivalence.lean, which shows the typed-graph interpretation of the MLP equals this spec's forward pass; NN/Proofs/Models/Mlp.lean and NN/Tests/Runtime/Floats/TorchLeanSpecMlpEquivCheck.lean exercise it further.

def Examples.mlpSpec {α : Type} [TorchLean.Storage α] [Context α] {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec α inDim hidDim) (l2 : Spec.LinearSpec α hidDim outDim) : Spec.Module.Chain α [inDim] [outDim]

A 2-layer MLP as a Spec.Module.Chain:

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} [TorchLean.Storage α] [Context α] {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec α inDim hidDim) (l2 : Spec.LinearSpec α hidDim outDim) : Spec.Module.Chain α [inDim] [outDim]

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} [TorchLean.Storage α] [Context α] {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec α inDim hidDim) (l2 : Spec.LinearSpec α hidDim outDim) (x : TorchLean.Tensor α [inDim]) : TorchLean.Tensor α [outDim]

Run the MLP forward on a single input vector.

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

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} [TorchLean.Storage α] [Context α] {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec α inDim hidDim) (l2 : Spec.LinearSpec α hidDim outDim) (x : TorchLean.Tensor α [inDim]) : (mlpSpec l1 l2).forward x = have z1 := Spec.linearSpec l1 x; have a1 := Activation.reluSpec z1; Spec.linearSpec l2 a1

The composed Spec.Module.Chain forward equals the hand-written Linear → ReLU → Linear computation.

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

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} [TorchLean.Storage α] [Context α] {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec α inDim hidDim) (l2 : Spec.LinearSpec α hidDim outDim) (x : TorchLean.Tensor α [inDim]) (dLdy : TorchLean.Tensor α [outDim]) : TorchLean.Tensor α [inDim]

Composed backward of the MLP using the OpSpec chain.