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*) and the proved IRExec.ForwardGraph bridge.

Assumptions:

• Node ids index g.nodes consistently; getNode 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.tensorToPyFlat {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

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.

• bufferFull (attr : String) : ConstBinding
• paramFull (attr : String) : ConstBinding

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

@[instance_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.convKernelAttr (id : ℕ) : String

Attribute name for a convolution kernel tensor in the parameter store.

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

Attribute name for a convolution bias tensor in the parameter store.

def Export.IRPyTorch.natTensorToPyTuple {n : ℕ} (v : TorchLean.Tensor ℕ [n]) : String

Render a fixed-length natural-number vector as a Python tuple.

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

Attribute name for a BatchNorm scale tensor (self.batchnorm_<id>_gamma).

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

Attribute name for a BatchNorm bias tensor (self.batchnorm_<id>_beta).

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

Attribute name for a BatchNorm running mean tensor (self.batchnorm_<id>_mean).

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

Attribute name for a BatchNorm running variance tensor (self.batchnorm_<id>_var).

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