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:

• imports / support provide the shared Python prelude.
• base is the reusable class skeleton for the example exporters.
• emit is the simplest end-to-end exporter for a Spec.Module.Chain.
• checkpoints emits Python state-dictionary save/load helpers. Callers supply any checkpoint-to-JSON conversion needed by the import path.

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.boolLiteral (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.

Common boilerplate fragments

Many exporters emit the same small pieces of Python: @property metadata and a get_model_info dictionary. These fragments are shared by the example exporters and the general IR exporter.

def Export.PyTorch.metadata (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.shapeLiteral (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.tupleLiteral {n : ℕ} (dimensions : TorchLean.Tensor ℕ [n]) : String

Render tensor-valued dimensions as a Python tuple, including the empty and singleton cases.

def Export.PyTorch.floatLiteral (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.tensorLiteral {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.imports : String

Standard imports used by the generated Python snippets.

def Export.PyTorch.support : 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.base (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.checkpoints : String

Emit Python helpers for saving/loading state-dict checkpoints.

def Export.PyTorch.checks : String

Emit Python helpers for validating exported models.

def Export.PyTorch.emit {α : Type} [TorchLean.Storage α] {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. The scalar storage is inherited from the supplied chain; emission only reads its layer metadata.