torch.export / FX Graph JSON Import

This module is the Lean half of the PyTorch-module import pipeline.

PyTorch has two different artifacts that people often blur together:

• state_dict: tensor values keyed by names; great for weights, but it does not describe the architecture.
• torch.export / FX graph: a captured tensor program; this is the piece we need before TorchLean can run native semantics, verification, or proof-oriented analyses.

The runtime bridge therefore uses a small, explicit JSON format:

{
  "format": "torchlean.ir.v2",
  "input_ids": [0],
  "output_ids": [1],
  "nodes": [
    {"id": 0, "kind": "input", "parents": [], "shape": [1, 4]},
    {"id": 1, "kind": "relu", "parents": [0], "shape": [1, 4]}
  ]
}

The parser below is deliberately conservative. It accepts only the current NN.IR.OpKind subset and then runs TorchLean's executable graph validators (checkWellFormed and checkShapes). That gives downstream tools a useful guarantee: if parseGraph succeeds, the resulting graph is structurally acyclic, id-disciplined, arity-correct, and shape-consistent according to the shared IR inference rules.

Op kind strings are not spelled here. They come from the shared table in NN.Runtime.PyTorch.Wire, which the Python-emitting exporters also use. The tag round trip is proved there.

The matching Python-side emitter lives in NN.Runtime.PyTorch.Export.TorchExport.

structure Import.PyTorch.TorchExport.MhaPayloadBinding : Type

Node ids whose payloads are synthesized while lowering one raw single-head MHA tuple node.

Keeping this binding explicit prevents payload parsing from depending on incidental node-number offsets or reverse-engineering the generated subgraph.

• rawId : ℕ

  Raw value-graph node that owns the serialized MHA parameters.

• queryId : ℕ

  Generated query projection node.

• keyId : ℕ

  Generated key projection node.

• valueId : ℕ

  Generated value projection node.

• scaleId : ℕ

  Generated attention-score scale constant node.

• outputId : ℕ

  Generated output projection node.

@[instance_reducible]

instance Import.PyTorch.TorchExport.instReprMhaPayloadBinding : Repr MhaPayloadBinding

def Import.PyTorch.TorchExport.instReprMhaPayloadBinding.repr : MhaPayloadBinding → ℕ → Std.Format

structure Import.PyTorch.TorchExport.CapturedGraph : Type

A captured PyTorch graph lowered into TorchLean IR plus its designated interface node ids.

• graph : NN.IR.Graph

  TorchLean's checked op-tagged graph.

• inputId : ℕ

  Designated runtime input node id.

• outputIds : Array ℕ

  Designated graph output node ids, in PyTorch return order.

• rawToTensor : Array (Option ℕ)

  Map from serialized FX value ids to lowered tensor IR ids.

• mhaPayloadBindings : Array MhaPayloadBinding

  Payload ownership for subgraphs synthesized from raw MHA tuple nodes.

def Import.PyTorch.TorchExport.instReprCapturedGraph.repr : CapturedGraph → ℕ → Std.Format

@[instance_reducible]

instance Import.PyTorch.TorchExport.instReprCapturedGraph : Repr CapturedGraph

inductive Import.PyTorch.TorchExport.Internal.ValueShape : Type

Shape metadata for a PyTorch/FX value before it has been lowered to TorchLean's tensor-only IR.

PyTorch FX nodes may produce non-tensor containers. The most common example is nn.MultiheadAttention, whose forward result is (attn_output, attn_weights). We keep this container structure in the import layer instead of treating every FX node as a tensor node. Only tensor-valued nodes can be lowered into NN.IR.Graph.

• tensor (shape : Spec.Shape) : ValueShape
• tuple (items : Array Spec.Shape) : ValueShape

def Import.PyTorch.TorchExport.Internal.instReprValueShape.repr : ValueShape → ℕ → Std.Format

@[instance_reducible]

instance Import.PyTorch.TorchExport.Internal.instReprValueShape : Repr ValueShape

@[instance_reducible]

instance Import.PyTorch.TorchExport.Internal.instBEqValueShape : BEq ValueShape

def Import.PyTorch.TorchExport.Internal.instBEqValueShape.beq : ValueShape → ValueShape → Bool

structure Import.PyTorch.TorchExport.Internal.CapturedValueNode : Type

One raw tensor or container value from a graph artifact.

• id : ℕ

  Raw FX node id from the JSON artifact.

• parents : Array ℕ

  Raw parent ids. These may refer to tensor values or tuple/container values.

• kind : String

  Stable TorchLean/PyTorch import tag, e.g. relu, matmul, or tuple_getitem.

• valueShape : ValueShape

  Tensor or tuple shape metadata.

• raw : StateDict

  Original object, retained so tensor nodes can be parsed through kind.

structure Import.PyTorch.TorchExport.Internal.CapturedValueGraph : Type

Captured PyTorch graph before tuple/container values are lowered away.

• nodes : Array CapturedValueNode

  Raw FX/value nodes.

• inputId : ℕ

  Raw designated input id.

• outputIds : Array ℕ

  Raw designated output ids.

Small JSON helpers

def Import.PyTorch.TorchExport.Internal.typeError {α : Type} (ctx expected : String) (j : Lean.Json) : Except String α

A uniform "wrong JSON shape" error, so import failures all read the same way.

def Import.PyTorch.TorchExport.Internal.nat (ctx : String) : Lean.Json → Except String ℕ

Interpret a JSON number as a natural number.

def Import.PyTorch.TorchExport.Internal.nat.integral (value : ℕ) : ℕ → Option ℕ

def Import.PyTorch.TorchExport.Internal.string (ctx : String) : Lean.Json → Except String String

Decode a JSON value as a string, reporting the importing context on mismatch.

def Import.PyTorch.TorchExport.Internal.array (ctx : String) : Lean.Json → Except String (Array Lean.Json)

Decode a JSON value as an array, reporting the importing context on mismatch.

def Import.PyTorch.TorchExport.Internal.object (ctx : String) : Lean.Json → Except String StateDict

Decode a JSON value as an object/state dictionary.

def Import.PyTorch.TorchExport.Internal.field (ctx key : String) (o : StateDict) : Except String Lean.Json

Read a required object field.

def Import.PyTorch.TorchExport.Internal.nats (ctx : String) (j : Lean.Json) : Except String (Array ℕ)

Parse a JSON array of natural numbers.

def Import.PyTorch.TorchExport.Internal.shape (ctx : String) (j : Lean.Json) : Except String Spec.Shape

Parse a shape encoded as a dimension list.

Examples:

• [] means scalar;
• [4] means Shape.dim 4 Shape.scalar;
• [2, 3] means Shape.dim 2 (Shape.dim 3 Shape.scalar).

def Import.PyTorch.TorchExport.Internal.natField (ctx key : String) (o : StateDict) : Except String ℕ

Read a natural-number field from a parsed Torch export JSON object.

def Import.PyTorch.TorchExport.Internal.boolField (ctx key : String) (o : StateDict) : Except String Bool

Read a Boolean field from a parsed Torch export JSON object.

def Import.PyTorch.TorchExport.Internal.shapeField (ctx key : String) (o : StateDict) : Except String Spec.Shape

Read a shape-valued field from a parsed Torch export JSON object.

def Import.PyTorch.TorchExport.Internal.natArrayField (ctx key : String) (o : StateDict) : Except String (Array ℕ)

Read a natural-number array field from a parsed Torch export JSON object.

def Import.PyTorch.TorchExport.Internal.floatField (ctx key : String) (o : StateDict) : Except String Float

Read a floating-point field used by numerical operator payloads.

def Import.PyTorch.TorchExport.Internal.tensorField (ctx key : String) (shape : Spec.Shape) (o : StateDict) : Except String (TorchLean.Tensor Float shape)

Read a shape-indexed floating-point tensor field.

def Import.PyTorch.TorchExport.Internal.natTensorField (ctx key : String) (rank : ℕ) (o : StateDict) : Except String (TorchLean.Tensor ℕ [rank])

Read fixed-length tensor metadata, rejecting inconsistent serialized geometry.

def Import.PyTorch.TorchExport.Internal.windowConfig (ctx : String) (o : StateDict) : Except String NN.IR.WindowConfig

Parse generic pooling window metadata.

def Import.PyTorch.TorchExport.Internal.shapes (ctx : String) (j : Lean.Json) : Except String (Array Spec.Shape)

Parse a JSON array whose elements are shape arrays.

def Import.PyTorch.TorchExport.Internal.valueShape (ctx : String) (o : StateDict) : Except String ValueShape

Parse the value-level shape metadata emitted by the Python bridge.

A node without value_kind is tensor-valued and has a required shape field. The explicit form uses value_kind = "tensor" or value_kind = "tuple" explicitly.

def Import.PyTorch.TorchExport.Internal.pointwise (ctx tagString : String) (o : StateDict) : Except String NN.IR.OpKind

Parse a pointwise operator from its kind and attribute field.

A GELU must state its approximate string. A division without a rounding_mode field, or with a JSON null, is true division, as in PyTorch.

def Import.PyTorch.TorchExport.Internal.kind (ctx : String) (outShape : Spec.Shape) (o : StateDict) : Except String NN.IR.OpKind

Parse a TorchLean IR op kind.

The schema uses a stable string tag plus op-specific scalar fields. The tag is resolved through Wire.parseTag?, which recognizes the fixed spellings. This function then reads the attributes each constructor needs. We avoid trying to parse raw PyTorch operator names here; the Python adapter is responsible for translating torch.ops.aten.* / FX targets into these TorchLean tags.

def Import.PyTorch.TorchExport.Internal.node (j : Lean.Json) : Except String CapturedValueNode

Parse one raw PyTorch/FX value node.

def Import.PyTorch.TorchExport.Internal.inputs (j : Lean.Json) : Except String (Array ℕ)

Read the ordered tensor interface: input_ids lists every runtime input in argument order.

def Import.PyTorch.TorchExport.Internal.graph (j : Lean.Json) : Except String CapturedValueGraph

Parse the value-level graph, retaining the first argument for single-input consumers.

def Import.PyTorch.TorchExport.Internal.get (vg : CapturedValueGraph) (id : ℕ) : Except String CapturedValueNode

Look up a raw value node by id.

def Import.PyTorch.TorchExport.Internal.check (vg : CapturedValueGraph) (inputIds : Array ℕ := #[vg.inputId]) : Except String Unit

Validate the value graph before raw ids are used as array indices during tensor lowering.

The tensor IR checker performs the corresponding validation after lowering. This earlier check is still necessary because tuple/container nodes do not enter the tensor IR, and because rawToTensor is indexed by raw node id while lowering is in progress.

def Import.PyTorch.TorchExport.Internal.lower (vg : CapturedValueGraph) (inputIds : Array ℕ := #[vg.inputId]) : Except String CapturedGraph

Lower a PyTorch/FX value graph to TorchLean's tensor-only IR.

Tuple/container nodes are preserved in CapturedValueGraph, but they do not become NN.IR.Nodes. TorchLean's verification and execution passes consume the clean tensor DAG, while the import layer can explain container-valued PyTorch failures without changing their semantics.

def Import.PyTorch.TorchExport.Internal.validate (cg : CapturedGraph) : Except String CapturedGraph

Validate a lowered graph and its designated interface, preserving check order.

def Import.PyTorch.TorchExport.parseGraph (j : Lean.Json) : Except String CapturedGraph

Parse and validate a captured PyTorch graph.

Success means:

• the JSON uses the TorchLean graph-artifact schema,
• every op is in the supported TorchLean IR subset,
• node ids are disciplined and topologically ordered,
• arities are valid, and
• declared output shapes match NN.IR.Infer.

def Import.PyTorch.TorchExport.parsePayload (j : Lean.Json) : Except String (NN.IR.Payload Float)

Parse serialized parameters into the node-keyed Float payload.

theorem Import.PyTorch.TorchExport.parseGraph_wellShaped {j : Lean.Json} {cg : CapturedGraph} (h : parseGraph j = Except.ok cg) : cg.graph.WellShaped

Successfully parsed graphs satisfy the IR shape checks.