PyTorch graph capture

This module generates a Python script that captures an nn.Module and writes its graph as torchlean.ir.v1 JSON. The Lean importer reads that artifact and checks it before evaluation:

PyTorch nn.Module
  --torch.export / FX capture-->
TorchLean graph JSON (`torchlean.ir.v1`)
  --NN.Runtime.PyTorch.Import.TorchExport.parseGraph-->
NN.IR.Graph

PyTorch handles capture and shape propagation. The generated script translates supported operations into NN.IR.OpKind and reports an error for unsupported operations or configurations. The importer then checks the graph's edges, shapes, and operation attributes.

Operator matching uses exact callable or ATen identities. For example, torch.log can become an IR .log node, while torch.log1p needs a different formula and is rejected. Similar names and matching output shapes are insufficient to identify the computation.

The script is assembled from one definition per Python section so each piece can be read and changed on its own. Every "kind" string it writes comes from NN.Runtime.PyTorch.Wire, the same table the Lean importer parses with.

References:

• torch.export: https://docs.pytorch.org/docs/stable/user_guide/torch_compiler/export.html
• torch.fx: https://docs.pytorch.org/docs/stable/fx.html

structure Export.PyTorch.TorchExport.GraphBridgeOptions : Type

Options for the generated PyTorch graph-capture script.

• functionName : String

  Name of the Python helper function emitted into the script.

• preferTorchExport : Bool

  If true, use torch.export.export first and fall back to FX symbolic tracing.

• includeDebugTargets : Bool

  If true, include raw PyTorch target strings in each node for debugging.

@[instance_reducible]

instance Export.PyTorch.TorchExport.instReprGraphBridgeOptions : Repr GraphBridgeOptions

def Export.PyTorch.TorchExport.instReprGraphBridgeOptions.repr : GraphBridgeOptions → ℕ → Std.Format

Python sections

Each section is an array of Python lines ending in one blank separator line. Indentation is applied here so the sections concatenate into a valid module.

def Export.PyTorch.TorchExport.bridgeImports : Array String

Imports and the artifact format marker.

def Export.PyTorch.TorchExport.shapeHelpers : Array String

Shape readers for FX node metadata plus the getitem/getattr target tests.

def Export.PyTorch.TorchExport.tupleKindSection : Array String

Kind entries for tuple-valued FX nodes; nn.MultiheadAttention carries its payload.

def Export.PyTorch.TorchExport.modelHelpers : Array String

Model loading, target naming, FX node reference collection, and spatial tuple parsing.

def Export.PyTorch.TorchExport.payloadHelpers : Array String

Parameter payload serialization for linear, attention, and convolution nodes.

def Export.PyTorch.TorchExport.axisHelpers : Array String

Flatten-versus-reshape selection and axis normalization.

def Export.PyTorch.TorchExport.lowerKindHeader : Array String

Opening of _lower_kind: target name, argument, and axis extraction shared by all rules.

def Export.PyTorch.TorchExport.lowerModuleRules : Array String

_lower_kind rules for call_module nodes, keyed on the nn.Module subclass.

def Export.PyTorch.TorchExport.lowerElementwiseRules : Array String

_lower_kind rules for elementwise aten/functional targets that carry no payload.

def Export.PyTorch.TorchExport.lowerShapeRules : Array String

_lower_kind rules for reductions, softmax, and shape operations.

def Export.PyTorch.TorchExport.lowerFunctionalLayerRules : Array String

_lower_kind rules for functional layers with payloads: layer norm, linear, batch norm.

def Export.PyTorch.TorchExport.lowerFunctionalSpatialRules : Array String

_lower_kind rules for functional convolution and pooling, then the tuple/fallback cases.

def Export.PyTorch.TorchExport.lowerKindSection : Array String

The complete _lower_kind function.

def Export.PyTorch.TorchExport.lowerabilitySection : Array String

Check captured operations before dropping unused values.

FX records x.relu_(); return x with an unused result for the ReLU call. Skipping that call before checking its contract would erase the update to x. Unsupported calls therefore fail even when their results have no users.

def Export.PyTorch.TorchExport.captureSection (options : GraphBridgeOptions) : Array String

Graph capture: torch.export when preferred and lowerable, otherwise FX symbolic tracing.

def Export.PyTorch.TorchExport.exportFunctionSection (options : GraphBridgeOptions) : Array String

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

def Export.PyTorch.TorchExport.mainSection (options : GraphBridgeOptions) : Array String

Command-line main for the generated script.

def Export.PyTorch.TorchExport.generateGraphBridgeScript (options : GraphBridgeOptions := { }) : String

Emit a Python script that captures a PyTorch module and writes TorchLean graph JSON.

The generated script expects a Python file containing a zero-argument model constructor or class:

python export_torchlean_graph.py my_model.py MyModel out_graph.json --example-shape 1,4

The first implementation target is the shared IR subset: elementwise ops, matmul, reductions, reshape/permute/flatten/concat, softmax, layernorm, and simple pooling/conv metadata when PyTorch exposes enough static arguments. Linear layers can appear either as aten.linear or as lower-level matmul/add depending on PyTorch's graph capture.