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 that intentionally exclude .mse_loss.
• dValsOfCtx_*: typed-context to IR-array bridge lemmas.
• denoteAllState_* helpers: semantic equivalence bridges between compiled state and IR denotation tables.

Implementation notes

• We keep this module deliberately stable: it is shared infrastructure, so 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 avoids duplicated brittle simp chains across correctness modules.
• 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. The best long-term direction is fewer ad-hoc simp scripts in the op proofs and more named lemmas with clear contracts.

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.

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 is a small convenience lemma used throughout the compiled-correctness proofs to close impossible branches where compilation would have thrown an error message.

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 α] [Inhabited α] {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 α] [Inhabited α] {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⟩.

theorem Runtime.Autograd.Compiled.evalAt_matmul_mm_ok {α : Type} [Context α] [Inhabited α] [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 α] [Inhabited α] [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.

theorem Runtime.Autograd.Compiled.evalAt_reduceSum_ok {α : Type} [Context α] [Inhabited α] [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 : NN.IR.Graph.mkValidAxis? 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 ⋯))

NN.IR.Graph.evalAt for a .reduceSum axis node, specialized to 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_reduceMean_ok {α : Type} [Context α] [Inhabited α] [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 : NN.IR.Graph.mkValidAxis? 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.