MLP PyTorch Reference Export

PyTorch code generator for the MLP round-trip reference model.

The generated Python mirrors the common nn.Linear → ReLU → nn.Linear pattern. We also support embedding explicit weights into a state_dict-shaped dictionary for round-trip and regression checks.

inductive Export.PyTorch.MLP.KeyStyle : Type

How to name state_dict keys when exporting weights.

• linear : KeyStyle

  Keys like fc1.weight / fc2.bias (matches PyTorch nn.Linear modules).

• sequential : KeyStyle

  Keys like layers.0.weight / layers.2.bias (common when exporting nn.Sequential).

@[instance_reducible]

instance Export.PyTorch.MLP.instDecidableEqKeyStyle : DecidableEq KeyStyle

@[instance_reducible]

instance Export.PyTorch.MLP.instReprKeyStyle : Repr KeyStyle

def Export.PyTorch.MLP.instReprKeyStyle.repr : KeyStyle → ℕ → Std.Format

def Export.PyTorch.MLP.key (style : KeyStyle) (layer : ℕ) (field : String) : String

State-dictionary key for a zero-based linear-layer index and parameter field. Sequential models place a ReLU between linear layers, so their indices are doubled.

def Export.PyTorch.MLP.body (inputWidth hiddenWidth outputWidth : ℕ) (className : String) (softmax : Bool := false) : Array String

Emit the Python class body for a Linear → ReLU → Linear MLP, optionally ending in softmax.

This returns lines (not a single string) so callers can splice it into larger scripts.

def Export.PyTorch.MLP.source (inputWidth hiddenWidth outputWidth : ℕ) (className : String := "MLP") (softmax : Bool := false) : String

Render a standalone MLP class, optionally applying softmax after its output layer.

def Export.PyTorch.MLP.weights {inputWidth hiddenWidth outputWidth : ℕ} (inputWeight : TorchLean.Tensor Float [hiddenWidth, inputWidth]) (inputBias : TorchLean.Tensor Float [hiddenWidth]) (outputWeight : TorchLean.Tensor Float [outputWidth, hiddenWidth]) (outputBias : TorchLean.Tensor Float [outputWidth]) (className : String := "MLP") (keyStyle : KeyStyle := KeyStyle.linear) (softmax : Bool := false) : String

Generate Python code for an MLP plus helper functions that embed concrete weights.

The output contains a get_mlp_state_dict function that returns a PyTorch-shaped dictionary (state_dict). Its load_mlp_weights helper normalizes either key convention to the generated class's fc1/fc2 layers before calling model.load_state_dict(...). Set softmax := true to apply softmax after the final linear layer; parameter keys are unchanged.

def Export.PyTorch.MLP.script {inputWidth hiddenWidth outputWidth : ℕ} (className : String := "MLP") : String

Generate a complete Python script for MLP examples.

This includes:

• a base MLP class,
• a Softmax variant,
• shared checkpoint and test utilities,
• and convenience helpers for construction and parameter counting.