GraphSpec MLP Example

This file contains the smallest GraphSpec architecture example:

Linear(input, hidden) → ReLU → Linear(hidden, output).

This does not duplicate TorchLean's executable MLP helper. Application code uses TorchLean.nn.mlp; this file keeps only the proof-oriented graph description. The point here is narrower and proof-oriented:

• show the sequential Chain DSL in its simplest useful form;
• make the parameter ABI visible in the type;
• provide a stable target for GraphSpec equivalence and deterministic-init proofs.

Because this is a pure sequential chain, it is authored with Chain and >>>. The general LowerToDAG.Chain.toDAGModelZeroInit conversion also makes it available to DAG tooling.

@[reducible, inline]

abbrev NN.GraphSpec.Models.mlpParams (inputWidth hiddenWidth outputWidth : ℕ) : List Spec.Shape

Parameter ABI for the 2-layer MLP: (W₁, b₁, W₂, b₂).

def NN.GraphSpec.Models.mlp (inputWidth hiddenWidth outputWidth : ℕ) : Chain (mlpParams inputWidth hiddenWidth outputWidth) [inputWidth] [outputWidth]

2-layer MLP: Linear(input, hidden) → ReLU → Linear(hidden, output).

Notice how the parameter interface is explicit in the type:

• the first linear layer contributes tensors W₁ : Tensor α [hiddenWidth, inputWidth] and b₁ : Tensor α [hiddenWidth],
• the second linear layer contributes tensors W₂ : Tensor α [outputWidth, hiddenWidth] and b₂ : Tensor α [outputWidth],
• and ReLU contributes no parameters.

So the overall parameter list is exactly: [[hiddenWidth, inputWidth], [hiddenWidth], [outputWidth, hiddenWidth], [outputWidth]].

Example Usage

You can build a simple classifier head by appending a softmax:

def g (inputWidth hiddenWidth outputWidth : Nat) :
    Chain
      [[hiddenWidth, inputWidth], [hiddenWidth], [outputWidth, hiddenWidth], [outputWidth]]
      [inputWidth] [outputWidth] :=
  Models.mlp inputWidth hiddenWidth outputWidth >>> Chain.softmax [outputWidth] 0

Then:

• Interp.spec (g …) maps a parameter pack and input tensor to an output tensor;
• Chain.toProgram (g …) is an executable TorchLean Program with arguments params ++ [input].