Export Core

PyTorch code generation helpers.

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

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 instead emits plain Python source primarily for readability and for round-trip tests.

Reading map:

• generatePyTorchImports / generatePyTorchHelperModules 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 SpecChain.
• generateCompletePyTorchExport combines the codegen pieces into a single script.
• 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

structure Export.PyTorch.PyTorchExportMetadata (α : Type) (s t : Spec.Shape) : Type

Metadata for a generated PyTorch model snippet.

Most exporters in NN/Runtime/PyTorch/Export/* produce a Python string as their final artifact. We keep a small structured record alongside that string so demos can report shapes/layer counts and decide whether weights were embedded.

• modelName : String

  Human-friendly class/model name used in the emitted Python.

• inputShape : Spec.Shape

  Expected input shape (TorchLean spec Shape).

• outputShape : Spec.Shape

  Expected output shape (TorchLean spec Shape).

• layerCount : ℕ

  Count of primitive layers/ops in the model (exporter-specific).

• operationTypes : List String

  A list of operation tags used for lightweight reporting in demos.

• hasWeights : Bool

  Whether the exporter embedded a state_dict literal in the emitted Python.

• pytorchCode : String

  The emitted Python source (usually a full script or class definition).

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

Join a list of lines with newline separators.

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

Indent a line by n spaces.

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

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

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

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

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

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

def Export.PyTorch.indent8 (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 demo exporters and the more general IR exporter.

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

Emit a standard get_model_info method used by most TorchLean PyTorch demo 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.generateGetModelInfoMethodLines.renderItems : List (String × String) → List String

def Export.PyTorch.shapeDims : Spec.Shape → List ℕ

Flatten a Shape into a list of dimension sizes (outermost-first).

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} : ModSpec.SpecChain α s t → ℕ

Count the number of primitive layers in a SpecChain.

def Export.PyTorch.tensor1DToList {n : ℕ} (t : Spec.Tensor Float (Spec.Shape.dim n Spec.Shape.scalar)) : List Float

Convert a 1D float tensor into a Lean list (outermost dimension order).

def Export.PyTorch.tensor1DToPy {n : ℕ} (t : Spec.Tensor Float (Spec.Shape.dim n Spec.Shape.scalar)) : String

Render a 1D float tensor as a Python list literal (e.g. [1.0, 2.0]).

def Export.PyTorch.tensor2DToPy {rows cols : ℕ} (t : Spec.Tensor Float (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) : String

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

def Export.PyTorch.tensor2DToPyT {rows cols : ℕ} (t : Spec.Tensor Float (Spec.Shape.dim rows (Spec.Shape.dim cols Spec.Shape.scalar))) : 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.tensor4DToPy {a b c d : ℕ} (t : Spec.Tensor Float (Spec.Shape.dim a (Spec.Shape.dim b (Spec.Shape.dim c (Spec.Shape.dim d Spec.Shape.scalar))))) : String

Render a 4D float tensor as a Python nested-list literal.

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

Best-effort conversion of a float tensor to a Python list literal.

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

def Export.PyTorch.generatePyTorchImports : String

Standard imports used by the generated Python snippets.

def Export.PyTorch.generatePyTorchHelperModules : String

Small helper modules used by some toPyTorch 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 lightweight checkpoints.

def Export.PyTorch.generateTestingUtils : String

Emit lightweight Python helpers for validating exported models.

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

Generate a complete nn.Sequential-based Python module for a SpecChain.

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

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

Like generateBasePyTorchModule, but also include shared weight/testing helpers.

def Export.PyTorch.exportGeneralModel {α : Type} {s t : Spec.Shape} (chain : ModSpec.SpecChain α s t) (className : String := "GeneralModel") : PyTorchExportMetadata α s t

Export a SpecChain to PyTorch and bundle the result in a metadata record.

This is the "one call" entrypoint used by some demos.