Export Core

PyTorch code generation helpers.

This module defines shared string-building utilities used by the PyTorch bridge and round-trip examples. It emits readable Python nn.Module code (optionally with weights embedded) and centralizes the common prelude used by NN.Runtime.PyTorch.Export.{MLP,CNN,Transformer}.

Design note (PyTorch export APIs, for context only):

PyTorch also has graph capture / serialization mechanisms such as ONNX export and torch.export. Those APIs produce IR-like artifacts intended for execution in other runtimes. TorchLean's exporter in this folder emits auditable Python source for parity checks and round-trip tests.

The public helpers are organized as follows:

• generatePyTorchImports / generatePyTorchSupportDefinitions provide the shared Python prelude.
• generateBasePyTorchModule is the reusable class skeleton for the example exporters.
• generatePyTorchModule is the simplest end-to-end exporter for a Spec.Module.Chain.
• NN.Runtime.PyTorch.Export.StateDict is the general checkpoint-to-JSON adapter for users who already have PyTorch weights.

References

• PyTorch ONNX export: https://pytorch.org/docs/stable/onnx.html
• PyTorch torch.export: https://pytorch.org/docs/stable/export.html

def Export.PyTorch.joinLines (xs : Array String) : String

Join an array of lines with newline separators.

def Export.PyTorch.pyBool (b : Bool) : String

Render a Lean Bool as the corresponding Python literal.

def Export.PyTorch.indent (n : ℕ) (s : String) : String

Indent a line by n spaces.

def Export.PyTorch.indentTwo (s : String) : String

Indent a line by 2 spaces (common for Python).

def Export.PyTorch.indentFour (s : String) : String

Indent a line by 4 spaces (common for Python block bodies).

def Export.PyTorch.indentSix (s : String) : String

Indent a line by 6 spaces (used in nested Python blocks).

def Export.PyTorch.indentEight (s : String) : String

Indent a line by 8 spaces (used for nested nn.Sequential strings).

Common boilerplate fragments

Many exporters emit the same small pieces of Python: @property metadata and a get_model_info dictionary. Shared boilerplate keeps the hand-written example exporters and the more general IR exporter.

def Export.PyTorch.generateGetModelInfoMethodLines (modelName : String) (extraFields : Array (String × String) := #[]) : Array String

Emit a standard get_model_info method used by most TorchLean PyTorch example modules.

extraFields are inserted after the "model_name" entry. Each element is (key, valueExpr) where valueExpr is emitted verbatim as Python code (e.g. "self.input_shape" or "self.hidden_dim"). This is meant as a formatting helper only; it does not validate Python syntax.

def Export.PyTorch.shapeToPyTupleString (s : Spec.Shape) : String

Render a Shape as a Python tuple literal.

Examples:

• .scalar becomes "()",
• a 1D shape becomes "(n,)" (note the trailing comma),
• higher-rank shapes become "(d0, d1, ...)".

def Export.PyTorch.countLayers {α : Type} {s t : Spec.Shape} : Spec.Module.Chain α s t → ℕ

Count the number of primitive layers in a Spec.Module.Chain.

def Export.PyTorch.floatToPyString (value : Float) : String

Render a Python float expression preserving every finite binary64 value and signed zero.

Short decimal strings are retained only when they parse back to the original bits. Otherwise the expression uses Python's built-in float.fromhex with the exact integer significand and binary exponent. Infinities and NaN use explicit Python constructors; NaN payload bits are not serialized.

def Export.PyTorch.tensorToPyString {s : Spec.Shape} (t : TorchLean.Tensor Float s) : String

Convert a float tensor to a Python list literal without rounding its binary64 elements.

This is a simple recursive printer used for examples and small regression tests; it is not intended to be fast.

def Export.PyTorch.transposedMatrixTensorToPy {rows cols : ℕ} (t : TorchLean.Tensor Float [rows, cols]) : String

Render the transpose of a 2D float tensor as a Python nested-list literal.

TorchLean's matrix-valued specs often follow the mathematical convention where a feature matrix W has shape (in, out) and is applied as X * W. PyTorch stores nn.Linear weights as (out, in) and applies them as X @ W.T + b. This helper prints a TorchLean matrix in the transposed orientation expected by PyTorch.

def Export.PyTorch.generatePyTorchImports : String

Standard imports used by the generated Python snippets.

def Export.PyTorch.generatePyTorchSupportDefinitions : String

Small helper modules used by some pythonExpr strings in the Lean specs.

These are small, dependency-free Python utilities (selectors, wrappers, a compact attention helper) used so the generated model classes stay short and readable.

def Export.PyTorch.generateBasePyTorchModule (className docstring : String) : String

Generate a generic base nn.Module class skeleton.

This is used by exporters that want a "real" class with an explicit _initialize_layers hook, instead of the simpler nn.Sequential emitter.

def Export.PyTorch.generateWeightLoadingUtils : String

Emit Python helpers for saving/loading state dictionaries and JSON checkpoints.

def Export.PyTorch.generateTestingUtils : String

Emit Python helpers for validating exported models.

def Export.PyTorch.generatePyTorchModule {α : Type} {s t : Spec.Shape} (chain : Spec.Module.Chain α s t) (className : String := "ExportedModel") : String

Generate a complete nn.Sequential-based Python module for a Spec.Module.Chain.

This is the simplest exporter: we extract an array of (opName, pythonLayerString) pairs and drop them into an nn.Sequential(...) in a new class.