Transformer PyTorch Reference Export

PyTorch code generator for the Transformer encoder round-trip reference model.

This file produces a readable Python nn.Module implementation that follows the usual PyTorch structure (MHA + residual + LayerNorm + FFN). In the TorchLean repo we mostly use this as a round-trip companion: generate a reference implementation, train/tweak in Python if needed, and optionally export parameters back to Lean via JSON in the importer modules.

def Export.PyTorch.Transformer.classSource (sequenceLength modelWidth headCount feedForwardWidth layerCount : ℕ) (className : String := "TransformerEncoder") : String

Render a small Transformer encoder as a Python nn.Module class definition.

This produces readable "reference PyTorch" code (MultiHeadAttention + residual + LayerNorm + FFN), useful for round-trip examples.

def Export.PyTorch.Transformer.withParameters (sequenceLength modelWidth headCount feedForwardWidth : ℕ) (queryWeight keyWeight valueWeight outputWeight : TorchLean.Tensor Float [modelWidth, modelWidth]) (feedForwardInputWeight : TorchLean.Tensor Float [feedForwardWidth, modelWidth]) (feedForwardOutputWeight : TorchLean.Tensor Float [modelWidth, feedForwardWidth]) (feedForwardInputBias : TorchLean.Tensor Float [feedForwardWidth]) (feedForwardOutputBias norm1Scale norm1Bias norm2Scale norm2Bias : TorchLean.Tensor Float [modelWidth]) (className : String := "TransformerEncoder") : String

Generate a single-layer Transformer encoder module with an embedded state_dict initializer.

This is meant for round-trip examples where parameters are loaded from TorchLean tensors.

TorchLean attention projections use mathematical (input, output) orientation and are transposed for PyTorch. Feed-forward layers already use PyTorch's (output, input) orientation and are emitted unchanged.