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

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

@[instance_reducible]

instance Export.PyTorch.MLP.instReprKeyStyle : Repr KeyStyle

def Export.PyTorch.MLP.firstWeightKey : KeyStyle → String

Key name for the first layer's weight tensor in a PyTorch state_dict.

def Export.PyTorch.MLP.firstBiasKey : KeyStyle → String

Key name for the first layer's bias tensor in a PyTorch state_dict.

def Export.PyTorch.MLP.secondWeightKey : KeyStyle → String

Key name for the second layer's weight tensor in a PyTorch state_dict.

def Export.PyTorch.MLP.secondBiasKey : KeyStyle → String

Key name for the second layer's bias tensor in a PyTorch state_dict.

def Export.PyTorch.MLP.classLines (inputWidth hiddenWidth outputWidth : ℕ) (className : String) : Array String

Emit the Python class body for a basic Linear → ReLU → Linear MLP.

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

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

Render a standalone Python file containing an nn.Module MLP class.

def Export.PyTorch.MLP.withParameters {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) : 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(...).

def Export.PyTorch.MLP.softmaxClassLines {inputWidth hiddenWidth outputWidth : ℕ} (className : String) : Array String

Render a line-based MLP class with a terminal softmax.

def Export.PyTorch.MLP.completeSource {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 helper modules from NN/Runtime/PyTorch/Export/Core.lean,
• and convenience helpers for construction and parameter counting.