ONNX Graph Adapter

This module emits a Python-side ONNX adapter for TorchLean graph import.

The adapter does not give ONNX a separate Lean semantics. It reads an ONNX model, lowers the supported static graph fragment to the same torchlean.ir.v2 JSON used by the torch.export/FX bridge, and then the Lean side should call Import.PyTorch.TorchExport.parseGraph.

That boundary is intentional: ONNX parsing, protobuf handling, and shape inference stay outside Lean; TorchLean accepts only the small checked graph artifact.

Every "kind" string written by the adapter comes from NN.Runtime.PyTorch.Wire, the table the Lean importer parses with.

structure Export.PyTorch.ONNX.Options : Type

Options for the generated ONNX-to-TorchLean-IR adapter script.

• functionName : String

  Name of the Python helper function emitted into the script.

• includeDebugTargets : Bool

  Include ONNX node names/op types in each node for debugging.

@[instance_reducible]

instance Export.PyTorch.ONNX.instReprOptions : Repr Options

def Export.PyTorch.ONNX.instReprOptions.repr : Options → ℕ → Std.Format

Python sections

Each section is an array of Python lines ending in one blank separator line.

def Export.PyTorch.ONNX.Internal.imports : Array String

Imports and the artifact format marker.

def Export.PyTorch.ONNX.Internal.shapes : Array String

Static tensor metadata, arithmetic preconditions, and attribute readers.

def Export.PyTorch.ONNX.Internal.rule (opType : String) (tag : NN.IR.OpTag) : String

One-line _lower_node rule mapping an ONNX op type to a payload-free kind.

def Export.PyTorch.ONNX.Internal.nodes : Array String

_lower_node: single-node lowerings keyed on the ONNX op type.

def Export.PyTorch.ONNX.Internal.reductions : Array String

_reduce_kind: ReduceSum/ReduceMean with ONNX defaults. Axes come from the axes attribute (before opset 13 for ReduceSum, before opset 18 for ReduceMean) or from a constant second input. Missing or empty axes reduce every dimension unless noop_with_empty_axes is set, and keepdims defaults to 1.

def Export.PyTorch.ONNX.Internal.softmax : Array String

_lower_softmax_coerced: Softmax before opset 13, whose default axis is 1 and which normalizes over the input coerced to 2-D [prod(d < axis), prod(d >= axis)].

def Export.PyTorch.ONNX.Internal.unary : List NN.IR.OpTag

Kinds whose single parent is the first ONNX input.

def Export.PyTorch.ONNX.Internal.binary : List NN.IR.OpTag

Kinds that take exactly two ONNX inputs as parents.

def Export.PyTorch.ONNX.Internal.parents : Array String

_ir_parent_names: select the ONNX inputs that become IR parents for a lowered kind.

def Export.PyTorch.ONNX.Internal.payload : Array String

_conv_payload: the dense weight and the bias of a Conv whose parameters are initializers.

Grouped kernels [O, I/group, k...] are expanded to the dense [O, I, k...] layout the IR conv payload uses, with zeros outside each group's block, as the torch.export adapter does. A missing bias becomes zeros.

def Export.PyTorch.ONNX.Internal.convolution : Array String

_lower_conv: channels-first Conv with explicit padding, optionally unwrapping a batch.

def Export.PyTorch.ONNX.Internal.gemm : Array String

_lower_gemm: Gemm as matmul plus an optional broadcast bias add.

def Export.PyTorch.ONNX.Internal.batchNorm : Array String

_lower_batchnorm: inference-mode BatchNormalization expanded into elementwise nodes.

Each per-channel [C] tensor is reshaped to [C, 1, ..., 1] before broadcasting, because IR broadcastTo aligns shapes from the right and the channel axis of [N, C, ...] is axis 1.

def Export.PyTorch.ONNX.Internal.entrypoint (options : Options) : Array String

The exported entry point: walk the ONNX graph and write the artifact.

def Export.PyTorch.ONNX.Internal.cli (options : Options) : Array String

Command-line main for the generated script.

def Export.PyTorch.ONNX.script (options : Options := { }) : String

Emit a Python script that lowers a conservative ONNX fragment to torchlean.ir.v2.

Supported first-pass ops are tensor ops that map to current NN.IR.OpKind: elementwise arithmetic/activations, MatMul, ReduceSum, ReduceMean, Softmax, Reshape, Flatten, Concat, Transpose, Gemm, inference-style BatchNormalization, and arbitrary-rank channels-first Conv. Flatten becomes the 2-D reshape ONNX specifies, and Softmax before opset 13 is lowered through the same 2-D coercion. Graph initializers become constant nodes whose values Import.PyTorch.TorchExport.parsePayload reads back. A Conv node carries its dense weight and bias in the payload, so its weight and optional bias must be initializers. Only homogeneous floating-point data from the standard ONNX operator domain is accepted. Integer shape and reduction-axis initializers remain metadata, not floating-point operands. Square roots require structurally nonnegative operands; the same check applies inside BatchNorm. NaN or infinite values are rejected because the JSON artifact cannot represent them.