Jaxpr graph import

NN/Runtime/JAX/jaxpr.py lowers captured JAX computations to the shared graph schema. The artifact lists the positional tensor inputs in input_ids. We reuse the existing graph parser and shape checker. parseProgram imports per-node types and explicit casts; parseBinaryPayload loads exact binary words. parseGraph retains the scalar-polymorphic path and rejects mixed precision.

parsePayload loads finite constants into native Float for numerical comparisons. That evaluator uses binary64 arithmetic; it does not reproduce a captured binary16, bfloat16, or binary32 rounding schedule. The graph can also be studied under its real-valued semantics, with the assumptions required by each operation.

The adapter and XLA are outside the proof boundary. Accepted graphs are shape checked, not certified translations of their source programs.

def Import.JAX.Internal.checkFrontend (artifact : Lean.Json) : Except String Unit

Check the frontend marker before using the shared IR readers.

def Import.JAX.parseProgram (artifact : Lean.Json) : Except String NN.IR.Program

Import a Jaxpr with per-node precision and explicit binary conversions.

def Import.JAX.parseBinaryPayload (artifact : Lean.Json) : Except String NN.IR.Interpreter.binary.Payload

Load binary constants without converting their values through native binary64.

def Import.JAX.Internal.checkPrecision (artifact : Lean.Json) : Except String Unit

Check that dtype metadata does not claim a cast the homogeneous IR cannot represent.

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

Import a Jaxpr artifact through the shared structural and shape validators.

def Import.JAX.parsePayload (artifact : Lean.Json) : Except String (NN.IR.Payload Float)

Load constants for the native binary64 diagnostic evaluator, not source-dtype replay.

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

An accepted Jaxpr artifact satisfies the same shape invariant as other imported graphs.