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.MLPPyTorch.WeightKeyStyle : Type

How to name state_dict keys when exporting weights.

• linear : WeightKeyStyle

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

• sequential : WeightKeyStyle

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

@[implicit_reducible]

instance Export.MLPPyTorch.instDecidableEqWeightKeyStyle : DecidableEq WeightKeyStyle

@[implicit_reducible]

instance Export.MLPPyTorch.instReprWeightKeyStyle : Repr WeightKeyStyle

def Export.MLPPyTorch.instReprWeightKeyStyle.repr : WeightKeyStyle → ℕ → Std.Format

def Export.MLPPyTorch.w1Key : WeightKeyStyle → String

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

def Export.MLPPyTorch.b1Key : WeightKeyStyle → String

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

def Export.MLPPyTorch.w2Key : WeightKeyStyle → String

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

def Export.MLPPyTorch.b2Key : WeightKeyStyle → String

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

structure Export.MLPPyTorch.MLPExportMetadata (α : Type) (inDim hidDim outDim : ℕ) : Type

Metadata for an exported MLP.

This is a small “record of facts” about the generated Python: shapes, dimensions, and optionally embedded weights for round-trip tests.

• modelName : String

  Python class name used in the generated module.

• inputDim : ℕ

  Input feature dimension.

• hiddenDim : ℕ

  Hidden layer dimension.

• outputDim : ℕ

  Output feature dimension.

• hasWeights : Bool

  Whether the generated Python should embed concrete weights.

• weights : Option (Spec.Tensor Float (Spec.Shape.dim hidDim (Spec.Shape.dim inDim Spec.Shape.scalar)) × Spec.Tensor Float (Spec.Shape.dim hidDim Spec.Shape.scalar) × Spec.Tensor Float (Spec.Shape.dim outDim (Spec.Shape.dim hidDim Spec.Shape.scalar)) × Spec.Tensor Float (Spec.Shape.dim outDim Spec.Shape.scalar))

  Optional two-layer MLP parameter payload: weights and biases for both linear layers.

def Export.MLPPyTorch.generateMLPPyTorchClassLines (inDim hidDim outDim : ℕ) (className : String) : List 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.MLPPyTorch.generateMLPPyTorchClass (inDim hidDim outDim : ℕ) (className : String := "MLP") : String

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

def Export.MLPPyTorch.generateMLPWithWeights {inDim hidDim outDim : ℕ} (w1 : Spec.Tensor Float (Spec.Shape.dim hidDim (Spec.Shape.dim inDim Spec.Shape.scalar))) (b1 : Spec.Tensor Float (Spec.Shape.dim hidDim Spec.Shape.scalar)) (w2 : Spec.Tensor Float (Spec.Shape.dim outDim (Spec.Shape.dim hidDim Spec.Shape.scalar))) (b2 : Spec.Tensor Float (Spec.Shape.dim outDim Spec.Shape.scalar)) (className : String := "MLP") (keyStyle : WeightKeyStyle := WeightKeyStyle.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) and a load_mlp_weights helper that calls model.load_state_dict(...).

def Export.MLPPyTorch.generateMLPWithSoftmax {inDim hidDim outDim : ℕ} (className : String := "MLPWithSoftmax") : String

Generate a Python MLP class that ends with a Softmax (classification convenience).

This mirrors common compact PyTorch model code; full training pipelines usually use logits + a combined loss (e.g. CrossEntropyLoss) instead of an explicit softmax.

def Export.MLPPyTorch.generateMLPWithSoftmaxLines {inDim hidDim outDim : ℕ} (className : String) : List String

Line-based version of generateMLPWithSoftmax for script composition.

def Export.MLPPyTorch.exportMLPFromSpecChain {α : Type} {inDim hidDim outDim : ℕ} (_chain : ModSpec.SpecChain α (Spec.Shape.dim inDim Spec.Shape.scalar) (Spec.Shape.dim outDim Spec.Shape.scalar)) (className : String := "ExportedMLP") : MLPExportMetadata α inDim hidDim outDim

Export metadata for an MLP described as a SpecChain.

The chain is accepted to keep the API aligned with SpecChain, while this exporter produces metadata from the explicit dimensions supplied by the type parameters.

def Export.MLPPyTorch.exportMLPWithWeights {α : Type} {inDim hidDim outDim : ℕ} (_chain : ModSpec.SpecChain α (Spec.Shape.dim inDim Spec.Shape.scalar) (Spec.Shape.dim outDim Spec.Shape.scalar)) (w1 : Spec.Tensor Float (Spec.Shape.dim hidDim (Spec.Shape.dim inDim Spec.Shape.scalar))) (b1 : Spec.Tensor Float (Spec.Shape.dim hidDim Spec.Shape.scalar)) (w2 : Spec.Tensor Float (Spec.Shape.dim outDim (Spec.Shape.dim hidDim Spec.Shape.scalar))) (b2 : Spec.Tensor Float (Spec.Shape.dim outDim Spec.Shape.scalar)) (className : String := "ExportedMLP") : MLPExportMetadata α inDim hidDim outDim

Like exportMLPFromSpecChain, but include explicit weights in the metadata record.

def Export.MLPPyTorch.generateCompleteMLPExport {inDim hidDim outDim : ℕ} (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.