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.v1",
  "input_id": 0,
  "output_id": 4,
  "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.

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

structure Import.PyTorch.TorchExport.CapturedGraph : Type

A captured PyTorch graph lowered into TorchLean IR plus the designated input/output node ids.

• graph : NN.IR.Graph

  TorchLean's checked op-tagged graph.

• inputId : ℕ

  Designated runtime input node id.

• outputId : ℕ

  Designated graph output node id.

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

@[implicit_reducible]

instance Import.PyTorch.TorchExport.instReprCapturedGraph : Repr CapturedGraph

inductive Import.PyTorch.TorchExport.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 : List Spec.Shape) : ValueShape

@[implicit_reducible]

instance Import.PyTorch.TorchExport.instReprValueShape : Repr ValueShape

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

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

@[implicit_reducible]

instance Import.PyTorch.TorchExport.instBEqValueShape : BEq ValueShape

structure Import.PyTorch.TorchExport.CapturedValueNode : Type

One raw PyTorch/FX value node from torchlean.ir.v1 JSON.

• id : ℕ

  Raw FX node id from the JSON artifact.

• parents : List ℕ

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

structure Import.PyTorch.TorchExport.CapturedValueGraph : Type

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

• nodes : Array CapturedValueNode

  Raw FX/value nodes.

• inputId : ℕ

  Raw designated input id.

• outputId : ℕ

  Raw designated output id.

Small JSON helpers

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

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

Interpret a JSON number as a natural number.

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

Interpret a JSON string.

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

Interpret a JSON array.

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

Interpret a JSON object.

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

Read a required object field.

def Import.PyTorch.TorchExport.field? (key : String) (o : StateDict) : Option Lean.Json

Read an optional object field.

def Import.PyTorch.TorchExport.parseNatList (ctx : String) (j : Lean.Json) : Except String (List ℕ)

Parse a JSON array of natural numbers.

def Import.PyTorch.TorchExport.shapeOfDims (dims : List ℕ) : Spec.Shape

Convert a list of dimensions into TorchLean's nested Shape representation.

def Import.PyTorch.TorchExport.parseShape (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.parseParents (ctx : String) (j : Lean.Json) : Except String (List ℕ)

Parse a node's parent id list.

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

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

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

def Import.PyTorch.TorchExport.natListField (ctx key : String) (o : StateDict) : Except String (List ℕ)

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

Parse a JSON array whose elements are shape arrays.

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

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

torchlean.ir.v1 artifacts without value_kind are interpreted as tensor-valued nodes with a required shape field. The explicit form uses value_kind = "tensor" or value_kind = "tuple" explicitly.

def Import.PyTorch.TorchExport.parseOpKind (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. We intentionally 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.parseNode (j : Lean.Json) : Except String NN.IR.Node

Parse one node object from the graph JSON format.

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

Parse one raw PyTorch/FX value node.

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

Parse the graph object into the PyTorch/FX value-level graph, before tensor lowering.

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

Look up a raw value node by id.

def Import.PyTorch.TorchExport.lowerValueGraph (vg : CapturedValueGraph) : 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.parseGraphUnchecked (j : Lean.Json) : Except String CapturedGraph

Parse the graph object and lower PyTorch/FX values to the tensor IR without validation.

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.

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

Guarantee exposed by the parser: a successfully parsed graph is well-shaped.

This is intentionally small but important. It is the theorem downstream verification/export code can quote when it receives a graph artifact through this importer.