Common

Internal helper lemmas for NN.Runtime.Autograd.Compiled.IRExec.Correctness.

These lemmas relate the typed runtime context (TList) to the untyped IR value table (Array DVal), and provide small “building block” correctness steps that are reused across the per-op proofs.

Reading map:

• dValsOfCtx* lemmas: relate the typed context produced by GraphData.eval to an untyped Array (NN.IR.DVal α) (this is what the IR evaluator uses).
• denoteAllState* lemmas: package the compiled evaluator (ExecGraphData.denoteAll) in the form expected by IR-style semantic equivalence proofs.

These lemmas are infrastructure: they should not encode op-specific logic. Per-op correctness files (Matmul/Pool2d/LayerNorm/MSELoss) should depend on this module and not re-prove these bridges.

Main definitions

• throw_bind_ne_ok: eliminates impossible success branches after throw.
• NoMSELoss: side condition for semantic equivalence theorems over fragments that exclude .mse_loss.
• NoRawLog: side condition for theorem statements that do not yet carry the positivity precondition required by raw real log.
• NoConcat: side condition for the concat branch not yet connected to end-to-end preservation.
• dValsOfCtx_*: typed-context to IR-array bridge lemmas.
• denoteAllState_* helpers: semantic equivalence bridges between compiled state and IR denotation tables.

Implementation notes

• This module is shared infrastructure: predictable proof contracts matter more than clever proof tricks.
• Many lemmas here are proof-irrelevance/indexing bridges; these are repetitive but they remove a lot of friction from op-specific proofs.
• Collecting these utilities in one place gives op-specific correctness modules shared rewrite and indexing lemmas instead of repeated local proof scripts.
• These files can build slowly because they connect two representations at once: typed TList contexts on the compiled side and dynamically shaped DVal arrays on the IR side. Most of the cost is not arithmetic; it is Lean checking that shape casts, array indices, and proof-irrelevant casts line up exactly.
• When the same proof pattern appears in multiple operator files, prefer a named lemma with a clear contract over another local simp script.

Tags

correctness, infrastructure, tlist, dval, bridge-lemmas

theorem Runtime.Autograd.Compiled.throw_eq_error {β : Type} (msg : String) : throw msg = Except.error msg

throw msg in the Except monad is .error msg.

Shared side conditions

These predicates describe the exact fragment covered by a theorem. Keeping them in Common lets the per-op lemmas, the semantic equivalence proof, and the chapter index refer to the same public contract without import cycles.

def Runtime.Autograd.Compiled.NoMSELoss (g : NN.IR.Graph) : Prop

Core semantic equivalence side condition: the IR graph contains no .mse_loss nodes.

The compiled runtime has a correct .mse_loss step lemma in Correctness.Ops.Loss; the existing end-to-end semantic equivalence theorem keeps this condition so its branch proof stays small and predictable.

def Runtime.Autograd.Compiled.NoRawLog (g : NN.IR.Graph) : Prop

Core semantic equivalence side condition: the IR graph contains no raw .log nodes.

Raw real logarithm is only mathematical on positive inputs. The executable closure produced by IRExec.buildFrom is intentionally pure, while the IR denotation can report an Except.error for nonpositive inputs. Until the compiler theorem carries a per-node positivity/domain-validity predicate, the end-to-end semantic-equivalence theorem excludes raw .log. Use the total epsilon-protected safe_log operator when a graph needs unconditional execution.

def Runtime.Autograd.Compiled.NoConcat (g : NN.IR.Graph) : Prop

Core semantic equivalence side condition: the IR graph contains no .concat nodes.

IRExec.buildFrom executes concat, including nonzero-axis permutation lowering, but the current end-to-end proof does not yet relate that branch to NN.IR.Graph.evalAt. Keeping this limitation as an explicit theorem hypothesis avoids presenting runtime support as a proved compiler fragment.

theorem Runtime.Autograd.Compiled.throw_bind_ne_ok {β γ : Type} {msg : String} {k : β → Except String γ} {v : γ} (h : (do let y ← throw msg k y) = Except.ok v) : False

If a do-chain begins with throw, it cannot produce an .ok result.

This lemma is used throughout the compiled-correctness proofs to close impossible branches where compilation would have thrown an error message.

theorem Runtime.Autograd.Compiled.exceptUnit_two_bind_first_ok {β : Type} {e₁ e₂ : Except String Unit} {next : Except String β} {v : β} (h : (do e₁ e₂ next) = Except.ok v) : e₁ = Except.ok ()

If two unit guards and a tail computation return .ok, then the first guard succeeded.

theorem Runtime.Autograd.Compiled.exceptUnit_two_bind_second_ok {β : Type} {e₁ e₂ : Except String Unit} {next : Except String β} {v : β} (h : (do e₁ e₂ next) = Except.ok v) : e₂ = Except.ok ()

If two unit guards and a tail computation return .ok, then the second guard succeeded.

theorem Runtime.Autograd.Compiled.exceptUnit_two_bind_tail_ok {β : Type} {e₁ e₂ : Except String Unit} {next : Except String β} {v : β} (h : (do e₁ e₂ next) = Except.ok v) : next = Except.ok v

If two unit guards and a tail computation return .ok, then the tail returned .ok.

theorem Runtime.Autograd.Compiled.array_getElem_proof_irrel {β : Type} (xs : Array β) (i : ℕ) (h₁ h₂ : i < xs.size) : xs[i] = xs[i]

Array indexing is proof-irrelevant.

This is a small technical lemma: in Lean, xs[i]'h carries a proof h : i < xs.size. Different proofs should not change the value returned by indexing.

theorem Runtime.Autograd.Compiled.array_getElem!_eq_getElem {β : Type} [Inhabited β] (xs : Array β) (i : ℕ) (h : i < xs.size) : xs[i]! = xs[i]

Relate xs[i]! (defaulting lookup) and xs[i]'h (bounded lookup) when the index is in-bounds.

This is a small bridge lemma used throughout the IR/runtime context comparison proofs.

@[simp]

theorem Runtime.Autograd.Compiled.dValsOfCtx_cast {α : Type} [Context α] {ss₁ ss₂ : List Spec.Shape} (h : ss₁ = ss₂) (ctx : Proofs.Autograd.Algebra.TList α ss₁) : IRExec.dValsOfCtx (Proofs.Autograd.Algebra.TList.cast h ctx) = IRExec.dValsOfCtx ctx

dValsOfCtx ignores type-level casts of the underlying TList.

GraphData.eval introduces a definitional cast when extending contexts; this lemma lets us erase it before reasoning about the corresponding Array of DVals.

@[simp]

theorem Runtime.Autograd.Compiled.dValsOfCtx_snoc {α : Type} [Context α] {ss : List Spec.Shape} {τ : Spec.Shape} (ctx : Proofs.Autograd.Algebra.TList α ss) (t : Spec.Tensor α τ) : IRExec.dValsOfCtx (ctx.snoc t) = (IRExec.dValsOfCtx ctx).push (NN.IR.DVal.mk τ t)

dValsOfCtx for a snoc’d context corresponds to Array.push of the appended tensor.

theorem Runtime.Autograd.Compiled.dValsOfCtx_getElem! {α : Type} [Context α] {ss : List Spec.Shape} (ctx : Proofs.Autograd.Algebra.TList α ss) (i : Fin ss.length) : (IRExec.dValsOfCtx ctx)[↑i]! = NN.IR.DVal.mk (ss.get i) (ctx.get i)

Indexing dValsOfCtx agrees with indexing the underlying TList context.

This is the main bridge between the typed runtime context and the untyped IR value table.

theorem Runtime.Autograd.Compiled.dValsOfCtx_getIdx {α : Type} [Context α] {ss : List Spec.Shape} {s : Spec.Shape} (ctx : Proofs.Autograd.Algebra.TList α ss) (idx : Proofs.Autograd.Algebra.Idx ss s) : (IRExec.dValsOfCtx ctx)[↑idx.i]! = NN.IR.DVal.mk s (Proofs.Autograd.Algebra.getIdx ctx idx)

Indexing dValsOfCtx by a typed Idx agrees with getIdx on the underlying TList.

This packages dValsOfCtx_getElem! into the repository’s Idx wrapper.

@[simp]

theorem Runtime.Autograd.Compiled.Graph.expectShape_mk {α : Type} [Context α] [DecidableEq Spec.Shape] {s : Spec.Shape} (t : Spec.Tensor α s) : NN.IR.Graph.expectShape s (NN.IR.DVal.mk s t) = Except.ok t

Graph.expectShape succeeds on a DVal built with the same shape.

@[simp]

theorem Runtime.Autograd.Compiled.Graph.expectShape_sigma {α : Type} [Context α] [DecidableEq Spec.Shape] {s : Spec.Shape} (t : Spec.Tensor α s) : NN.IR.Graph.expectShape s ⟨s, t⟩ = Except.ok t

Same as Graph.expectShape_mk, but for the sigma-style constructor ⟨s, t⟩.

inductive Runtime.Autograd.Compiled.MatmulOperands (α : Type) [Context α] : Type

Well-typed operands for either matrix-matrix or batched matrix multiplication.

• matrices {α : Type} [Context α] {m n p : ℕ} (left : Spec.Tensor α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (right : Spec.Tensor α (Spec.Shape.dim n (Spec.Shape.dim p Spec.Shape.scalar))) : MatmulOperands α

  Two rank-two matrices with matching inner dimensions.

• batches {α : Type} [Context α] {batch m n p : ℕ} (left : Spec.Tensor α (Spec.Shape.dim batch (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar)))) (right : Spec.Tensor α (Spec.Shape.dim batch (Spec.Shape.dim n (Spec.Shape.dim p Spec.Shape.scalar)))) : MatmulOperands α

  Two rank-three matrix batches with matching batch and inner dimensions.

def Runtime.Autograd.Compiled.MatmulOperands.leftDVal {α : Type} [Context α] : MatmulOperands α → NN.IR.DVal α

Package the left typed operand as a dynamic IR value.

def Runtime.Autograd.Compiled.MatmulOperands.rightDVal {α : Type} [Context α] : MatmulOperands α → NN.IR.DVal α

Package the right typed operand as a dynamic IR value.

def Runtime.Autograd.Compiled.MatmulOperands.resultDVal {α : Type} [Context α] : MatmulOperands α → NN.IR.DVal α

Evaluate a well-typed operand pair with the matching tensor specification.

def Runtime.Autograd.Compiled.MatmulOperands.resultAtShape {α : Type} [Context α] (operands : MatmulOperands α) (outShape : Spec.Shape) (shapeEq : operands.resultDVal.shape = outShape) : NN.IR.DVal α

Retag the typed multiplication result with an equal graph-declared output shape.

theorem Runtime.Autograd.Compiled.evalAt_matmul_ok {α : Type} [Context α] [DecidableEq Spec.Shape] (g : NN.IR.Graph) (payload : NN.IR.Payload α) (input : NN.IR.DVal α) (vals : Array (NN.IR.DVal α)) (i : ℕ) (n : NN.IR.Node) (aId bId : ℕ) (operands : MatmulOperands α) (hN : g.getNode i = Except.ok n) (hk : n.kind = NN.IR.OpKind.matmul) (hp : n.parents = [aId, bId]) (hGetA : vals[aId]! = operands.leftDVal) (hGetB : vals[bId]! = operands.rightDVal) (hOut : operands.resultDVal.shape = n.outShape) : g.evalAt payload input vals i = Except.ok (operands.resultAtShape n.outShape hOut)

NN.IR.Graph.evalAt agrees with the typed matrix-multiplication specification for both supported operand ranks. Rank compatibility is carried by MatmulOperands, so the evaluator proof is stated once instead of duplicating its graph and parent-value reasoning for mm and bmm.

theorem Runtime.Autograd.Compiled.evalAt_matmul_mm_ok {α : Type} [Context α] [DecidableEq Spec.Shape] (g : NN.IR.Graph) (payload : NN.IR.Payload α) (input : NN.IR.DVal α) (vals : Array (NN.IR.DVal α)) (i : ℕ) (n : NN.IR.Node) (aId bId m nDim p : ℕ) (aT : Spec.Tensor α (Spec.Shape.dim m (Spec.Shape.dim nDim Spec.Shape.scalar))) (bT : Spec.Tensor α (Spec.Shape.dim nDim (Spec.Shape.dim p Spec.Shape.scalar))) (hN : g.getNode i = Except.ok n) (hk : n.kind = NN.IR.OpKind.matmul) (hp : n.parents = [aId, bId]) (hGetA : vals[aId]! = NN.IR.DVal.mk (Spec.Shape.dim m (Spec.Shape.dim nDim Spec.Shape.scalar)) aT) (hGetB : vals[bId]! = NN.IR.DVal.mk (Spec.Shape.dim nDim (Spec.Shape.dim p Spec.Shape.scalar)) bT) (hOut : Spec.Shape.dim m (Spec.Shape.dim p Spec.Shape.scalar) = n.outShape) : g.evalAt payload input vals i = Except.ok (NN.IR.DVal.mk n.outShape (hOut ▸ Spec.matMulSpec aT bT))

NN.IR.Graph.evalAt for a .matmul node specialized to 2D matrix multiply.

This is a proof-only helper that records the exact Spec.mat_mul_spec term produced by evalAt in the well-typed success case.

theorem Runtime.Autograd.Compiled.evalAt_matmul_bmm_ok {α : Type} [Context α] [DecidableEq Spec.Shape] (g : NN.IR.Graph) (payload : NN.IR.Payload α) (input : NN.IR.DVal α) (vals : Array (NN.IR.DVal α)) (i : ℕ) (n : NN.IR.Node) (aId bId batch m nDim p : ℕ) (aT : Spec.Tensor α (Spec.Shape.dim batch (Spec.Shape.dim m (Spec.Shape.dim nDim Spec.Shape.scalar)))) (bT : Spec.Tensor α (Spec.Shape.dim batch (Spec.Shape.dim nDim (Spec.Shape.dim p Spec.Shape.scalar)))) (hN : g.getNode i = Except.ok n) (hk : n.kind = NN.IR.OpKind.matmul) (hp : n.parents = [aId, bId]) (hGetA : vals[aId]! = NN.IR.DVal.mk (Spec.Shape.dim batch (Spec.Shape.dim m (Spec.Shape.dim nDim Spec.Shape.scalar))) aT) (hGetB : vals[bId]! = NN.IR.DVal.mk (Spec.Shape.dim batch (Spec.Shape.dim nDim (Spec.Shape.dim p Spec.Shape.scalar))) bT) (hOut : Spec.Shape.dim batch (Spec.Shape.dim m (Spec.Shape.dim p Spec.Shape.scalar)) = n.outShape) : g.evalAt payload input vals i = Except.ok (NN.IR.DVal.mk n.outShape (hOut ▸ aT.bmmSpec bT))

NN.IR.Graph.evalAt for a .matmul node specialized to batched matmul (bmm).

Like evalAt_matmul_mm_ok, this is used to relate the IR evaluator’s result to the compiled node’s forward closure during the semantic equivalence correctness proof.

inductive Runtime.Autograd.Compiled.AxisReductionKind : Type

The two axis reductions supported by compiled IR semantic equivalence.

• sum : AxisReductionKind
• mean : AxisReductionKind

def Runtime.Autograd.Compiled.AxisReductionKind.toOpKind (operation : AxisReductionKind) (axis : ℕ) : NN.IR.OpKind

Convert a compiled axis-reduction case to its IR operation kind.

def Runtime.Autograd.Compiled.AxisReductionKind.denote {β : Type} [Context β] {shape : Spec.Shape} (operation : AxisReductionKind) (axis : ℕ) (tensor : Spec.Tensor β shape) (axisValid : Spec.Shape.valid_axis axis shape) : Spec.Tensor β (Spec.Tensor.shapeAfterSum shape axis)

Typed denotation of a compiled axis-reduction case.

theorem Runtime.Autograd.Compiled.evalAt_axisReduction_ok {α : Type} [Context α] [DecidableEq Spec.Shape] (operation : AxisReductionKind) (g : NN.IR.Graph) (payload : NN.IR.Payload α) (input : NN.IR.DVal α) (vals : Array (NN.IR.DVal α)) (i : ℕ) (n : NN.IR.Node) (pId axis : ℕ) (s : Spec.Shape) (pT : Spec.Tensor α s) (hAxisPf : PLift (Spec.Shape.valid_axis axis s)) (hN : g.getNode i = Except.ok n) (hk : n.kind = operation.toOpKind axis) (hp : n.parents = [pId]) (hGet : vals[pId]! = NN.IR.DVal.mk s pT) (hAxis : Spec.Shape.validAxis? axis s = some hAxisPf) (hOut : Spec.Tensor.shapeAfterSum s axis = n.outShape) : g.evalAt payload input vals i = Except.ok (NN.IR.DVal.mk n.outShape (hOut ▸ operation.denote axis pT ⋯))

NN.IR.Graph.evalAt for either axis-reduction node in a well-typed success case.

This helper records the exact Tensor.reduceSum term produced by the IR evaluator once:

• the parent has the expected shape s,
• the axis validity check succeeds, and
• the node's declared outShape matches shapeAfterSum s axis.

The final cast to n.outShape comes from the evalAt "shape-tag normalization" step.

theorem Runtime.Autograd.Compiled.evalAt_reduceSum_ok {α : Type} [Context α] [DecidableEq Spec.Shape] (g : NN.IR.Graph) (payload : NN.IR.Payload α) (input : NN.IR.DVal α) (vals : Array (NN.IR.DVal α)) (i : ℕ) (n : NN.IR.Node) (pId axis : ℕ) (s : Spec.Shape) (pT : Spec.Tensor α s) (hAxisPf : PLift (Spec.Shape.valid_axis axis s)) (hN : g.getNode i = Except.ok n) (hk : n.kind = NN.IR.OpKind.reduceSum axis) (hp : n.parents = [pId]) (hGet : vals[pId]! = NN.IR.DVal.mk s pT) (hAxis : Spec.Shape.validAxis? axis s = some hAxisPf) (hOut : Spec.Tensor.shapeAfterSum s axis = n.outShape) : g.evalAt payload input vals i = Except.ok (NN.IR.DVal.mk n.outShape (hOut ▸ Spec.Tensor.reduceSum axis pT ⋯))

evalAt_axisReduction_ok specialized to summation.

theorem Runtime.Autograd.Compiled.evalAt_reduceMean_ok {α : Type} [Context α] [DecidableEq Spec.Shape] (g : NN.IR.Graph) (payload : NN.IR.Payload α) (input : NN.IR.DVal α) (vals : Array (NN.IR.DVal α)) (i : ℕ) (n : NN.IR.Node) (pId axis : ℕ) (s : Spec.Shape) (pT : Spec.Tensor α s) (hAxisPf : PLift (Spec.Shape.valid_axis axis s)) (hN : g.getNode i = Except.ok n) (hk : n.kind = NN.IR.OpKind.reduceMean axis) (hp : n.parents = [pId]) (hGet : vals[pId]! = NN.IR.DVal.mk s pT) (hAxis : Spec.Shape.validAxis? axis s = some hAxisPf) (hOut : Spec.Tensor.shapeAfterSum s axis = n.outShape) : g.evalAt payload input vals i = Except.ok (NN.IR.DVal.mk n.outShape (hOut ▸ Spec.Tensor.reduceMean axis pT ⋯))

NN.IR.Graph.evalAt for a .reduceMean axis node, specialized to a well-typed success case.

This is the mean analogue of evalAt_reduceSum_ok.

def Runtime.Autograd.Compiled.execOfState {α : Type} (inShape : Spec.Shape) (st : IRExec.State α inShape) : ExecGraphData α

Repackage a compiled State as an ExecGraphData so we can call its evaluator helpers.

def Runtime.Autograd.Compiled.denoteAllState {α : Type} [Context α] (inShape : Spec.Shape) (st : IRExec.State α inShape) (x : Spec.Tensor α inShape) : Array (NN.IR.DVal α)

Evaluate the compiled prefix state and convert its typed runtime context into an IR-style table.

theorem Runtime.Autograd.Compiled.denoteAllState_snoc {α : Type} [Context α] {inShape : Spec.Shape} {ss : List Spec.Shape} {τ : Spec.Shape} (gd : Proofs.Autograd.Algebra.GraphData α Unit [inShape] ss) (nodeData : Proofs.Autograd.Algebra.NodeData α Unit ([inShape] ++ ss) τ) (x : Spec.Tensor α inShape) : have st := ⟨ss, gd⟩; have st' := ⟨ss ++ [τ], gd.snoc nodeData⟩; denoteAllState inShape st' x = (denoteAllState inShape st x).push (NN.IR.DVal.mk τ (nodeData.forward (gd.eval (Proofs.Autograd.Algebra.TList.cons x Proofs.Autograd.Algebra.TList.nil) ()) ()))

denoteAllState commutes with extending the SSA graph by one node (GraphData.snoc).

This is the key step for proving that the compiler’s prefix-building loop stays in semantic equivalence with the IR denotation table.

theorem Runtime.Autograd.Compiled.mkIdx_ok_i_eq [DecidableEq Spec.Shape] {inShape : Spec.Shape} {ss : List Spec.Shape} {id : ℕ} {s : Spec.Shape} {idx : Proofs.Autograd.Algebra.Idx ([inShape] ++ ss) s} (h : IRExec.mkIdx inShape ss id s = Except.ok idx) : ↑idx.i = id

Build a typed runtime index (Idx) for a numeric IR parent id.

The compiled runtime context is typed by a list of shapes [inShape] ++ ss. mkIdx checks that:

• id is in bounds, and
• the context shape at that position matches the expected shape s.

theorem Runtime.Autograd.Compiled.denoteAllState_get_mkIdx {α : Type} [Context α] [DecidableEq Spec.Shape] {inShape : Spec.Shape} {ss : List Spec.Shape} (gd : Proofs.Autograd.Algebra.GraphData α Unit [inShape] ss) (x : Spec.Tensor α inShape) {pid : ℕ} {s : Spec.Shape} {idx : Proofs.Autograd.Algebra.Idx ([inShape] ++ ss) s} (hIdx : IRExec.mkIdx inShape ss pid s = Except.ok idx) : (denoteAllState inShape ⟨ss, gd⟩ x)[pid]! = NN.IR.DVal.mk s (Proofs.Autograd.Algebra.getIdx (gd.eval (Proofs.Autograd.Algebra.TList.cons x Proofs.Autograd.Algebra.TList.nil) ()) idx)

Lookup lemma: denoteAllState[..][pid]! agrees with getIdx when mkIdx pid s succeeds.

This is used when proving correctness of the per-node compiler step: we translate parent ids in the IR into typed indices into the compiled context.

theorem Runtime.Autograd.Compiled.buildFrom_denoteAllFrom_finish {α : Type} [Context α] [DecidableEq Spec.Shape] (g : NN.IR.Graph) (payload : NN.IR.Payload α) {inShape : Spec.Shape} {ss : List Spec.Shape} (i : ℕ) (x : Spec.Tensor α inShape) (hi : i < g.nodes.size) (τ : Spec.Shape) (nodeData : Proofs.Autograd.Algebra.NodeData α Unit ([inShape] ++ ss) τ) (st1 st' : IRExec.State α inShape) (ctx : Proofs.Autograd.Algebra.TList α ([inShape] ++ ss)) (vals0 : Array (NN.IR.DVal α)) (input : NN.IR.DVal α) (hTail : g.denoteAllFrom payload input (i + 1) (denoteAllState inShape st1 x) = Except.ok (denoteAllState inShape st' x)) (hEval : g.evalAt payload input vals0 i = Except.ok (NN.IR.DVal.mk τ (nodeData.forward ctx ()))) (hStep : denoteAllState inShape st1 x = vals0.push (NN.IR.DVal.mk τ (nodeData.forward ctx ()))) : g.denoteAllFrom payload input i vals0 = Except.ok (denoteAllState inShape st' x)

One-step finishing lemma for the buildFrom/denoteAllFrom semantic equivalence proof.

If we know:

• the tail recursion i+1 is correct (hTail),
• the IR evaluator step at i matches the compiled node’s forward (hEval), and
• the compiled table at i is the previous table plus the pushed node value (hStep), then denoteAllFrom at i returns the final compiled table.