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 demos: 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 "source of truth" semantics live on the Lean side (NN.IR.Graph.denote*, and the compiled runtime).

Assumptions:

• Node ids index g.nodes consistently (we sanity-check this in getNode).
• The ParamStore contains parameters for the node kinds that require them (linear/conv2d/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.tensorToPyFlat {s : Spec.Shape} (t : Spec.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 : List Float) : Bool

Return true iff every element of the list equals the first element (vacuously true on []).

def Export.IRPyTorch.broadcastRow2D? (seqLen embedDim : ℕ) (flat : List Float) : Option (List Float)

Detect whether a flattened (seqLen × embedDim) tensor is a broadcast of a single row.

If so, return that row. This is used to compress large broadcasted constants into a smaller learnable vector parameter in the emitted PyTorch code.

structure Export.IRPyTorch.Options : Type

Options controlling IR-to-PyTorch emission.

Most knobs here configure demo emission (training scaffold, dtype), and how to materialize constant nodes (buffers vs parameters, and whether to compress broadcasted 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 smoke-test in the emitted script.

• learnableConsts : Bool

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

• compressBroadcastParams : Bool

  If true, compress broadcasted 2D consts into a smaller learnable vector parameter.

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

@[implicit_reducible]

instance Export.IRPyTorch.instReprOptions : Repr Options

inductive Export.IRPyTorch.ConstBinding : Type

How an IR const node is represented in the emitted PyTorch module.

• bufferFull: a non-learnable register_buffer(...) tensor
• paramFull: a learnable nn.Parameter with the full tensor shape
• paramBroadcast2D: a learnable vector nn.Parameter that is expanded at runtime to a (seqLen × embedDim) tensor (compression for broadcasted 2D constants).

• bufferFull (attr : String) : ConstBinding
• paramFull (attr : String) : ConstBinding
• paramBroadcast2D (attr : String) (seqLen embedDim : ℕ) : ConstBinding

def Export.IRPyTorch.instReprConstBinding.repr : ConstBinding → ℕ → Std.Format

@[implicit_reducible]

instance Export.IRPyTorch.instReprConstBinding : Repr ConstBinding

@[reducible, inline]

abbrev Export.IRPyTorch.ConstBindings : Type

Map from IR node id to how its constant should be referenced in the PyTorch module.

def Export.IRPyTorch.ConstBinding.attr : ConstBinding → String

The Python attribute name used to reference a bound constant (self.<attr>).

def Export.IRPyTorch.constAttr (id : ℕ) : String

Default attribute name for a const node: self.const_<id>.

def Export.IRPyTorch.linearWAttr (id : ℕ) : String

Attribute name for a linear layer weight tensor in the ParamStore (self.linear_<id>_W).

def Export.IRPyTorch.linearBAttr (id : ℕ) : String

Attribute name for a linear layer bias tensor in the ParamStore (self.linear_<id>_b).

def Export.IRPyTorch.conv2dKernelAttr (id : ℕ) : String

Attribute name for a conv2d kernel tensor in the ParamStore (self.conv2d_<id>_kernel).

def Export.IRPyTorch.conv2dBiasAttr (id : ℕ) : String

Attribute name for a conv2d bias tensor in the ParamStore (self.conv2d_<id>_bias).

def Export.IRPyTorch.pyTensorFromFlat (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.getNode (g : NN.IR.Graph) (id : ℕ) : Except String NN.IR.Node

Retrieve a node from the graph, with a couple sanity checks.

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 : ℕ) (opts : 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 scaffold for export checks.