TorchLean

5.4. Crossing Lean's Boundary🔗

PyTorch can capture a model, an interval library can propose a numerical enclosure, and a CUDA kernel can compute a tensor. TorchLean records what Lean learns when each external program returns.

There are two main boundaries:

  • a subprocess exchanges files, standard output, or JSON with another executable;

  • an FFI call invokes a linked native symbol and may exchange opaque memory handles.

Both can be used with explicit validation and named trust boundaries. They have different failure modes and support different proof stories.

5.4.1. A Subprocess Is An Untrusted Producer🔗

The common process helper is small:

def runJsonStdoutChecked
    (ctx : String)
    (cmd : String)
    (args : Array String)
    (cwd : Option String := some ".") :
    IO Json := do
  let stdout ← runStdoutChecked ctx cmd args cwd
  match Json.parse stdout with
  | .ok value => pure value
  | .error message =>
      throw <| IO.userError
        s!"{ctx}: JSON parse error: {message}\nstdout:\n{stdout}"

runStdoutChecked starts the process, captures its streams, and rejects a nonzero exit code with the command, arguments, status, and standard error in the diagnostic. runJsonStdoutChecked then requires all of standard output to be one JSON document.

Suppose Python prints:

{"format":"torchlean.bound.v1","lower":0.12,"upper":0.31}

Successful parsing establishes only that this text is valid JSON. A useful checker must still:

  1. require the exact format string;

  2. require finite numeric fields;

  3. establish \mathrm{lower}\leq\mathrm{upper};

  4. connect the interval to a particular graph, payload, input set, and output;

  5. invoke a sound acceptance theorem.

Changing upper to 1e999 is a good boundary test. JSON syntax accepts the number, but converting it to a machine Float can produce infinity. Verification parsers therefore use expectFiniteFloatE or expectFieldFiniteFloatE, not the permissive float parser, whenever the certificate schema promises finite claims.

5.4.2. A Complete PyTorch Capture🔗

Run:

lake exe torchlean pytorch_export_check

The command asks Python and torch.export to capture several small nn.Modules, emits torchlean.ir.v1 JSON, parses each document in Lean, lowers supported values to NN.IR.Graph, and runs the graph validators.

The current run accepts:

TinyAddRelu                 nodes=3
TinyMLP                     nodes=4
TinyCheckpointMLP           nodes=4
TinyCNN                     nodes=4
TinyCNNHead                 nodes=6
TinyBatchNorm2d             nodes=2
TinyNormSoftmax             nodes=3
TinyTransformerishBlock     nodes=7
TinySelfAttentionOps        nodes=5
TinySingleHeadMHA           nodes=11

Every accepted line reports:

guarantee: WellShaped via parseGraph_wellShaped

The wording is precise. The parser theorem says that successful parsing yields a graph satisfying the executable well-shaped predicate. It does not say that Python captured the intended module or that every PyTorch operator has been translated correctly.

5.4.3. Inspect The Exchange Format🔗

A minimal captured graph looks like:

{
  "format": "torchlean.ir.v1",
  "input_id": 0,
  "output_id": 2,
  "nodes": [
    {"id": 0, "kind": "input", "parents": [], "shape": [1, 4]},
    {"id": 1, "kind": "relu",  "parents": [0], "shape": [1, 4]},
    {"id": 2, "kind": "sum",   "parents": [1], "shape": []}
  ]
}

The Python producer is responsible for translating raw FX or ATen names into stable TorchLean tags. The Lean parser does not accept arbitrary operator strings and guess their meaning. It recognizes a conservative list, parses operation-specific fields, constructs candidate nodes, and validates the result.

This division is intentional:

PyTorch module
   ↓ external capture, treated as an untrusted producer
FX/value graph
   ↓ explicit JSON schema
Lean parser
   ↓ checked structural lowering
NN.IR.Graph
   ↓ denotation / verifier / exporter
TorchLean analysis

The container-valued FX layer matters. nn.MultiheadAttention returns a tuple of attention output and attention weights. Treating every FX node as a tensor loses that fact. TorchLean first retains tuple shape metadata, then lowers only supported tensor projections. The producer is free to do a complicated capture; the Lean side believes only the smaller document that its parser and graph checks actually recognize.

5.4.4. Rejection Is Part Of The Interface🔗

The same command deliberately tries unsupported models.

torch.sort(...).values is rejected because the current producer has no lowering rule for that tuple-valued operation. A two-head nn.MultiheadAttention is captured as a tuple-valued node but then rejected with:

PyTorch graph import: node[2]:
`nn.MultiheadAttention` lowering supports only num_heads=1, got 2

This is useful behavior. Replacing the unsupported node by an identity or silently dropping the attention weights would produce a valid-looking graph for the wrong model.

Try adding torch.sort to the Python example without changing the Lean parser. The expected result is an explicit unsupported-operation failure, not partial import.

5.4.5. A State Dictionary Is Not A Graph🔗

A state_dict supplies named tensors. It does not describe data flow. A graph capture supplies operations and edges. It may refer to parameters without carrying the full checkpoint provenance.

Round-trip import therefore has two obligations:

Artifact

Checks

graph

operator subset, IDs, parents, shapes, attributes, input/output IDs

state dictionary

key mapping, tensor shape, flat length, layout, finite values

Run:

lake exe torchlean pytorch_roundtrip

This writes the generated MLP PyTorch artifacts under NN/Examples/Interop/PyTorch/MLP/. Open the generated model and parameter files together: neither one is a complete description of the executable network by itself.

5.4.6. From Import Success To Semantic Equality🔗

Three propositions are easy to conflate:

  1. Python produced a document and exited successfully.

  2. Lean parsed the document into a well-shaped supported graph.

  3. The graph's denotation equals the original PyTorch module on all admissible inputs.

The capture experiment and parseGraph_wellShaped establish the second proposition, conditional on the bytes received. The third needs a translation theorem for the supported producer or an explicit trust assumption about capture and lowering.

Parity tests on random inputs are excellent engineering evidence for that assumption. They can find transposes, axis errors, missing biases, and mask inversions. They do not universally quantify over every parameter and input.

5.4.7. ONNX Uses The Same Checked Ingress🔗

NN.Runtime.PyTorch.Export.ONNX emits a Python adapter for a conservative static-shape ONNX fragment. The producer handles elementwise operations, rank-two/limited batched matmul, reductions, softmax, reshape/flatten, concat, selected transposes, Gemm, inference BatchNorm, and ungrouped, undilated convolution with the layouts represented by the current IR. It writes the same torchlean.ir.v1 document consumed by Import.PyTorch.TorchExport.parseGraph.

ONNX protobuf parsing and shape inference therefore remain outside Lean. Lean checks the smaller IR artifact it receives. Graph initializers and payload-backed nodes still need a matching payload import; a well-shaped convolution node by itself does not contain authenticated weights.

5.4.8. Native FFI Calls Have A Different Risk🔗

A subprocess returns copied bytes owned by Lean after parsing. A native symbol can allocate, mutate, alias, or free memory behind an opaque Lean value.

The CUDA buffer boundary contains declarations such as:

@[extern "torchlean_cuda_buffer_of_float_array_with_token"]
opaque ofFloatArrayWithToken
    (values : @& FloatArray) (token : UInt32) : Buffer

@[extern "torchlean_cuda_buffer_to_float_array_io"]
opaque toFloatArrayIO (buffer : @& Buffer) : IO FloatArray

@[extern "torchlean_cuda_buffer_release_with_token"]
opaque releaseWithToken
    (buffer : @& Buffer) (token : UInt32) : UInt32

The type Buffer is opaque. Lean code cannot forge its internal pointer, but the C implementation still determines whether allocation, copying, finalization, and release are correct.

The shape-erased tape therefore does not trust a shape tag by itself. Before a buffer reaches a shape-indexed kernel, AnyBuffer.validate checks that every axis and the total element count fit the CUDA UInt32 ABI and that the native buffer has exactly that many elements. Convolution and pooling wrappers additionally reject zero strides and oversized dimension arrays, validate input, kernel, and output spatial products independently, and check the returned native output length. Runtime natural-number gather indices accept the full UInt32 range and reject larger values before conversion. These checks turn malformed runtime values into ordinary Lean errors; they do not prove that a correctly sized native kernel implements its specification.

The @& annotation marks a borrowed Lean object. Native code may inspect such an argument during the call but must not retain or decrement it. Native translation units repeat critical length and geometry checks because an FFI declaration is still a trusted boundary. The file-to-symbol map is documented in NN.Runtime.Autograd.Engine.Cuda.NativeSources.

5.4.9. Why The IO Token Exists🔗

An external declaration that appears pure may be common-subexpression eliminated or reused by Lean as if equal arguments always denote the same value. Allocation does not have that semantics: two uploads of the same host array should produce independently owned buffers.

The effectful wrapper obtains a changing monotonic-time token:

def ofFloatArrayIO (values : @& FloatArray) : IO Buffer := do
  let timestamp ← IO.monoNanosNow
  pure <| ofFloatArrayWithToken values (UInt32.ofNat timestamp)

The native function ignores the token numerically. Its presence makes each allocation depend on the surrounding IO sequence, preventing Lean from treating repeated uploads as one pure object.

Release has the same issue. releaseIO uses a token and is called only at an ownership boundary where no alias will be used again. The native finalizer remains safe after explicit release because the implementation nulls the pointer.

A stale alias after release is a memory-safety bug that the tensor's shape type cannot detect.

5.4.10. Workspaces And Backward🔗

Some forward kernels produce an output plus intermediates needed by their VJP. TorchLean represents that ownership as:

structure WithWorkspace where
  value : Buffer
  workspace : List Buffer := []

The tape node retains the workspace until backward has consumed it. Afterwards releaseWorkspaceThen or releaseAllThen threads cleanup through a retained result, so the native release cannot be erased as dead pure code.

For long training runs this prevents two forms of growth:

  • GPU allocations waiting for Lean external-object finalizers;

  • tape closures retaining workspaces after their VJP has run.

Allocator counters report live and peak bytes, allocation/free counts, wrapper counts, and device free memory. They are observability tools, not a proof that no native leak exists.

5.4.11. Profile The Host Path🔗

LeanProfiler records named IO spans and writes both a Perfetto-compatible trace and a JSON timing summary. It can sit around an existing TorchLean runner:

open LeanProfiler

def main : IO Unit :=
  profileFromEnvironment "training" do
    span "model.run" runModel

With LEAN_PROFILE=1, the trace shows nesting and order while the summary groups repeated spans and records process counters. The comparison command checks a new summary against a baseline with explicit absolute and relative tolerances.

A host timer does not automatically measure asynchronous device work. The TorchLean integration can wait for device completion before closing a span; that number is completion latency, not per-kernel CUDA or CUPTI time. The distinction matters when a short host call merely queues a long GPU operation.

5.4.12. External Oracles And Certificates🔗

TorchLean also calls arbitrary-precision or interval tools to propose bounds. A typical workflow is:

Lean writes an exact query
        ↓
Arb / python-flint computes an enclosure
        ↓
Lean parses rational midpoint-radius data
        ↓
a checker validates the enclosure or treats it as oracle evidence

Parsing exact rationals avoids an extra binary64 conversion at ingestion. It still does not verify the external interval algorithm. If the result is replayed by a proved checker, checker acceptance supports the checker's proposition. If it is only compared in a test, it remains an oracle.

The same producer/checker distinction applies to α,β-CROWN leaf dumps, PINN residual artifacts, ODE enclosures, and geometry certificates. Search can be large and external; the accepted schema and soundness theorem should stay small enough to audit.

5.4.13. Evidence Is Field-Specific🔗

A useful boundary report may say:

  • shape: proved by a typed constructor;

  • layout: guarded by length and row-major checks;

  • value: compared by a regression suite;

  • VJP: delegated to a trusted external provider;

  • provenance: native symbol torchlean_cuda_buffer_matmul.

One strong field does not upgrade the others. In particular, a proof of shape safety is not a proof of arithmetic, and a source-file link is provenance rather than evidence.

5.4.14. Reproduce And Break The Boundary🔗

The two most useful experiments are:

lake exe torchlean pytorch_export_check
lake -R -K cuda=true exe torchlean quickstart_mlp \
  --device cuda --steps 2 --show-backend

Then deliberately break one condition:

  1. add an unsupported PyTorch operation and observe import rejection;

  2. change a JSON shape and observe checkShapes reject it;

  3. request CUDA from a stub build and observe runtime availability rejection;

  4. pass a wrong-size Q buffer to the LibTorch SDPA test and observe the Lean/native guard reject it.

These failure paths are part of the contract. A boundary that reports only success is difficult to audit and easy to overstate.