IRPyTorch

IR → PyTorch code generation.

This module takes an op-tagged NN.IR.Graph plus a NN.MLTheory.CROWN.Graph.ParamStore payload and emits a standalone PyTorch nn.Module implementation as a Python source string.

What this is (and isn't):

• This is an extraction/convenience layer used in round-trip examples: run/train in Python, then optionally import weights back to Lean.
• This is not a formal proof of semantic equivalence between PyTorch execution and the Lean IR denotation. The reference semantics remain the Lean definitions (NN.IR.Graph.denote*).

Assumptions:

• Node ids index g.nodes consistently; node checks that invariant.
• The ParamStore contains parameters for the node kinds that require them (linear/convolution/etc.).
• Only a supported subset of IR node kinds is lowered.

Failure modes (reported as Except String):

• missing/malformed ParamStore entries,
• shape mismatches between graph nodes and parameters,
• unsupported node kinds.

PyTorch context (comments only):

PyTorch’s ONNX exporter and torch.export can capture graphs for execution in other runtimes. TorchLean’s emitter here prints readable Python that mirrors the IR, rather than producing an execution-focused serialized graph artifact.

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

Flatten a tensor and render it as a Python list literal (used for torch.tensor([...]).reshape(...)).

def Export.IRPyTorch.allEq (xs : Array Float) : Bool

Return true iff every array entry equals the first one, vacuously for an empty array.

structure Export.IRPyTorch.Options : Type

Options controlling IR-to-PyTorch emission.

Most knobs here configure example emission, dtype handling, and how to materialize constant nodes as buffers or parameters.

• className : String

  Class name to use in the emitted Python source.

• dtypeExpr : String

  Python expression used for the tensor dtype (e.g. "torch.float32").

• includeTrainingSkeleton : Bool

  If true, include a small training skeleton and runtime-check in the emitted script.

• learnableConsts : Bool

  If true, emit IR const nodes as nn.Parameter when appropriate (learnable).

@[instance_reducible]

instance Export.IRPyTorch.instReprOptions : Repr Options

def Export.IRPyTorch.instReprOptions.repr : Options → ℕ → Std.Format

@[reducible, inline]

abbrev Export.IRPyTorch.Bindings : Type

Map from IR constant node id to its Python attribute name, for both buffers and parameters.

def Export.IRPyTorch.attr (kind : String) (id : ℕ) (field : String := "") : String

Python attribute name for a node and optional parameter field. The spelling is shared by initialization and forward evaluation, including saved state keys.

def Export.IRPyTorch.tensor (flatList : String) (shape : Spec.Shape) (dtypeVar : String := "dtype") : String

Emit a torch.tensor([...]).reshape(...) expression from a flat Python list literal and a Shape.

def Export.IRPyTorch.node (g : NN.IR.Graph) (id : ℕ) : Except String NN.IR.Node

Retrieve a node from the graph and validate its id invariant.

This yields more actionable error messages than directly indexing into g.nodes.

def Export.IRPyTorch.emit (g : NN.IR.Graph) (ps : NN.MLTheory.CROWN.Graph.ParamStore Float) (inputId outputId : ℕ) (options : Options := { }) : Except String String

Emit a standalone PyTorch nn.Module class for an IR graph.

This is the main entrypoint for IR exporters: it bundles:

• imports,
• a class definition with parameters/buffers materialized from the ParamStore,
• a forward method implementing the IR, and
• (optionally) a compact training file for export checks.