PyTorch nn.Module → TorchLean IR Check

This executable example exercises the model-agnostic PyTorch graph bridge:

1. write the generated Python graph-capture adapter;
2. write several compact supported PyTorch models (MLP/CNN/CNN head/attention/MHA blocks);
3. capture them to torchlean.ir.v1 JSON;
4. parse and validate each JSON artifact back into NN.IR.Graph; and
5. exercise a real torch.save(model.state_dict()) checkpoint reload path; and
6. run deliberate unsupported-op cases to confirm the bridge reports unsupported semantics clearly.

The check is scoped around the interop boundary rather than PyTorch performance. PyTorch may produce an artifact, but TorchLean accepts it only after parsing and validating the explicit IR JSON.

Run:

lake exe torchlean pytorch_export_check

def NN.Examples.Interop.PyTorch.TorchExportCheck.usage : String

Command-line help for the PyTorch export bridge check.

def NN.Examples.Interop.PyTorch.TorchExportCheck.workDir : System.FilePath

def NN.Examples.Interop.PyTorch.TorchExportCheck.bridgePath : System.FilePath

def NN.Examples.Interop.PyTorch.TorchExportCheck.modelPath : System.FilePath

def NN.Examples.Interop.PyTorch.TorchExportCheck.checkpointPath : System.FilePath

def NN.Examples.Interop.PyTorch.TorchExportCheck.supportedModelSource : String

def NN.Examples.Interop.PyTorch.TorchExportCheck.runCapture (ctor : String) (outPath : System.FilePath) (shape : String) : IO String

def NN.Examples.Interop.PyTorch.TorchExportCheck.writeCheckpoint : IO Unit

Write a real PyTorch checkpoint used by TinyCheckpointMLP.

Python owns .pt loading/saving, and the graph bridge only sees the resulting initialized nn.Module. That mirrors the intended user workflow: load trained weights in Python, capture the model graph, then ask Lean to validate the exported IR.

def NN.Examples.Interop.PyTorch.TorchExportCheck.compactFailure (err : IO.Error) : String

Keep subprocess failures readable in the runtime-check output.

The full IO.userError still contains command, args, exit code, and stderr. The runtime check only prints the first informative lines so unsupported-op tests explain the boundary without hiding the signal.

def NN.Examples.Interop.PyTorch.TorchExportCheck.runSupportedCase (ctor shape : String) : IO Unit

Run one supported capture path and parse the resulting graph in Lean.

def NN.Examples.Interop.PyTorch.TorchExportCheck.runSupported : IO Unit

Run the supported capture paths and parse the resulting graphs in Lean.

def NN.Examples.Interop.PyTorch.TorchExportCheck.runPythonUnsupportedCase (ctor shape msg : String) : IO Unit

Run a deliberate unsupported op and require the Python bridge to reject it.

def NN.Examples.Interop.PyTorch.TorchExportCheck.runLeanUnsupportedCase (ctor shape msg : String) : IO Unit

Run a PyTorch graph that can be captured as a value graph but must not lower into the tensor IR yet.

This is the intended behavior for container-valued PyTorch ops: the JSON keeps the tuple, then Lean rejects tensor lowering with a semantic explanation instead of the Python bridge crashing or incorrectly treating the tuple producer as a tensor op.

def NN.Examples.Interop.PyTorch.TorchExportCheck.runUnsupported : IO Unit

Run deliberate unsupported ops and require the Python bridge to reject them.

def NN.Examples.Interop.PyTorch.TorchExportCheck.run : IO Unit

Main runtime-check body.

def NN.Examples.Interop.PyTorch.TorchExportCheck.main (args : List String) : IO UInt32

Entrypoint used by lake exe torchlean pytorch_export_check.