Common

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

These lemmas relate the typed runtime context (TorchLean.TensorPack) to the untyped IR value table (Array Spec.SomeTensor) and provide small building-block correctness steps that are reused across the per-op proofs.

The lemmas are grouped as follows:

• packedTensorsOfContext* lemmas: relate the typed context produced by ForwardData.eval to an untyped Array (Spec.SomeTensor α) (this is what the IR evaluator uses).
• denoteAllState* lemmas: package the IR forward evaluator (ForwardGraph.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/Pooling/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.
• NoRawLog: side condition for theorem statements that do not carry the positivity precondition required by raw real log.
• noRawLog_of_forall_mem: discharge NoRawLog from a node-array fact, so a concrete graph can close it with decide.
• packedTensorsOfContext_*: typed-context to IR-array bridge lemmas.
• evalAt_matmul_leading_ok, evalAt_axisReduction_ok: evalAt in well-typed success cases.
• denoteAllState_* helpers: semantic equivalence bridges between lowered 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 TorchLean.TensorPack contexts on the forward-graph side and dynamically shaped Spec.SomeTensor 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, tensorpack, dval, bridge-lemmas

@[simp]

theorem Runtime.Autograd.IRExec.unaryParentId_eq_ok_of_unaryParent_eq_some (i p : ℕ) (n : NN.IR.Node) (h : NN.IR.unaryParent? n.parents = some p) : NN.IR.Graph.unaryParentId i n = Except.ok p

The lowering and semantic evaluators agree on a successfully decoded unary parent.

@[simp]

theorem Runtime.Autograd.IRExec.binaryParentIds_eq_ok_of_binaryParents_eq_some (i a b : ℕ) (n : NN.IR.Node) (h : NN.IR.binaryParents? n.parents = some (a, b)) : NN.IR.Graph.binaryParentIds i n = Except.ok (a, b)

The lowering and semantic evaluators agree on successfully decoded binary parents.

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.IRExec.NoRawLog (g : NN.IR.Graph) : Prop

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

The scientific raw-log domain is positive inputs. The IR denotation reports Except.error for nonpositive values, while the pure lowered closure applies Tensor.logSpec to every input. The end-to-end semantic-equivalence theorem therefore excludes every .log node. Protecting its input with a positive clamp or a softplus-based construction can avoid the domain failure, but does not discharge this syntactic predicate; that graph needs a separate domain-aware argument.

theorem Runtime.Autograd.IRExec.noRawLog_of_forall_mem {g : NN.IR.Graph} (h : ∀ n ∈ g.nodes, n.kind ≠ NN.IR.OpKind.log) : NoRawLog g

Discharge NoRawLog from a statement about the node array.

NoRawLog quantifies over node ids through Graph.getNode, which is the right shape for the recursive lowering proof but an annoying shape for a caller holding a concrete graph. Every id that getNode accepts reads an element of g.nodes, so a fact about the array is enough, and for a literal graph the hypothesis is closed by decide:

theorem myGraphNoRawLog : NoRawLog myGraph :=
  noRawLog_of_forall_mem (by decide)

We added this because the alternative was for every user of denoteAll_eq_of_lowerToForwardGraph to redo the same getNode case analysis inline.

theorem Runtime.Autograd.IRExec.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 lowered-correctness proofs to close impossible branches where lowering would have thrown an error message.

theorem Runtime.Autograd.IRExec.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.IRExec.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.IRExec.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.IRExec.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.

@[simp]

theorem Runtime.Autograd.IRExec.packedTensorsOfContext_cast {α : Type} [TorchLean.Storage α] {ss₁ ss₂ : List Spec.Shape} (h : ss₁ = ss₂) (ctx : TorchLean.TensorPack α ss₁) : Internal.packedTensorsOfContext (TorchLean.TensorPack.cast h ctx) = Internal.packedTensorsOfContext ctx

packedTensorsOfContext ignores type-level casts of the underlying TorchLean.TensorPack.

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

@[simp]

theorem Runtime.Autograd.IRExec.packedTensorsOfContext_snoc {α : Type} [TorchLean.Storage α] {ss : List Spec.Shape} {τ : Spec.Shape} (ctx : TorchLean.TensorPack α ss) (t : TorchLean.Tensor α τ) : Internal.packedTensorsOfContext (ctx.snoc t) = (Internal.packedTensorsOfContext ctx).push (Spec.SomeTensor.ofTensor t)

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

theorem Runtime.Autograd.IRExec.packedTensorsOfContext_getElem? {α : Type} [TorchLean.Storage α] {ss : List Spec.Shape} (ctx : TorchLean.TensorPack α ss) (i : Fin ss.length) : (Internal.packedTensorsOfContext ctx)[↑i]? = some (Spec.SomeTensor.ofTensor (ctx.get i))

Optional lookup in packedTensorsOfContext agrees with indexing the underlying typed context.

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

theorem Runtime.Autograd.IRExec.packedTensorsOfContext_getIdx? {α : Type} [TorchLean.Storage α] {ss : List Spec.Shape} {s : Spec.Shape} (ctx : TorchLean.TensorPack α ss) (idx : Proofs.Idx ss s) : (Internal.packedTensorsOfContext ctx)[↑idx.i]? = some (Spec.SomeTensor.ofTensor (Proofs.getIdx ctx idx))

Optional lookup in packedTensorsOfContext by a typed Idx agrees with getIdx on the underlying TorchLean.TensorPack.

This packages packedTensorsOfContext_getElem? into the repository’s Idx wrapper.

@[simp]

theorem Runtime.Autograd.IRExec.Graph.expectShape_mk {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (t : TorchLean.Tensor α s) : NN.IR.Graph.expectShape s { shape := s, tensor := t } = Except.ok t

Graph.expectShape succeeds on a Spec.SomeTensor built with the same shape.

theorem Runtime.Autograd.IRExec.Graph.expectShape_mk_of_eq {α : Type} [TorchLean.Storage α] [Context α] {s t : Spec.Shape} (h : s = t) (x : TorchLean.Tensor α s) : NN.IR.Graph.expectShape t { shape := s, tensor := x } = Except.ok (x.castShape h)

Graph.expectShape transports a shape-erased tensor when its stored shape equals the requested one.

theorem Runtime.Autograd.IRExec.evalAt_matmul_leading_ok {α : Type} [TorchLean.Storage α] [Context α] (g : NN.IR.Graph) (payload : NN.IR.Payload α) (input : Spec.SomeTensor α) (vals : Array (Spec.SomeTensor α)) (i : ℕ) (n : NN.IR.Node) (aId bId : ℕ) (leadingRev : List ℕ) (rows inner cols : ℕ) (aT : TorchLean.Tensor α ((Spec.Shape.ofList leadingRev.reverse).concat [rows, inner])) (bT : TorchLean.Tensor α ((Spec.Shape.ofList leadingRev.reverse).concat [inner, cols])) (hN : g.getNode i = Except.ok n) (hk : n.kind = NN.IR.OpKind.matmul) (hp : NN.IR.binaryParents? n.parents = some (aId, bId)) (hGetA : vals[aId]? = some { shape := (Spec.Shape.ofList leadingRev.reverse).concat [rows, inner], tensor := aT }) (hGetB : vals[bId]? = some { shape := (Spec.Shape.ofList leadingRev.reverse).concat [inner, cols], tensor := bT }) (hOut : (Spec.Shape.ofList leadingRev.reverse).concat [rows, cols] = n.outShape) : g.evalAt payload input vals i = Except.ok { shape := n.outShape, tensor := hOut ▸ NN.IR.Graph.matmulLeading (Spec.Shape.ofList leadingRev.reverse) aT bT }

NN.IR.Graph.evalAt for a .matmul node whose parents share the leading shape Shape.ofList leadingRev.reverse and end in the matrix axes [rows, inner] and [inner, cols].

The leading shape is spelled through its reversed dimension list because that is how both the IR evaluator and lowerMatmul recover it from the parent shapes. The lemma records the exact NN.IR.Graph.matmulLeading term produced by the evaluator for any leading shape (plain matrices, one batch axis, or several batch axes).

inductive Runtime.Autograd.IRExec.AxisReductionKind : Type

The two axis reductions supported by lowered IR semantic equivalence.

• sum : AxisReductionKind
• mean : AxisReductionKind

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

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

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

Typed denotation of a lowered axis-reduction case.

theorem Runtime.Autograd.IRExec.evalAt_axisReduction_ok {α : Type} [TorchLean.Storage α] [Context α] (operation : AxisReductionKind) (g : NN.IR.Graph) (payload : NN.IR.Payload α) (input : Spec.SomeTensor α) (vals : Array (Spec.SomeTensor α)) (i : ℕ) (n : NN.IR.Node) (pId axis : ℕ) (s : Spec.Shape) (pT : TorchLean.Tensor α s) (hAxisPf : PLift (Spec.Shape.NonemptyAxis axis s)) (hN : g.getNode i = Except.ok n) (hk : n.kind = operation.toOpKind axis) (hp : NN.IR.unaryParent? n.parents = some pId) (hGet : vals[pId]? = some { shape := s, tensor := pT }) (hAxis : Spec.Shape.nonemptyAxis? axis s = some hAxisPf) (hOut : TorchLean.Tensor.shapeAfterSum s axis = n.outShape) : g.evalAt payload input vals i = Except.ok { shape := n.outShape, tensor := 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.IRExec.evalAt_reduceSum_ok {α : Type} [TorchLean.Storage α] [Context α] (g : NN.IR.Graph) (payload : NN.IR.Payload α) (input : Spec.SomeTensor α) (vals : Array (Spec.SomeTensor α)) (i : ℕ) (n : NN.IR.Node) (pId axis : ℕ) (s : Spec.Shape) (pT : TorchLean.Tensor α s) (hAxisPf : PLift (Spec.Shape.NonemptyAxis axis s)) (hN : g.getNode i = Except.ok n) (hk : n.kind = NN.IR.OpKind.reduceSum axis) (hp : NN.IR.unaryParent? n.parents = some pId) (hGet : vals[pId]? = some { shape := s, tensor := pT }) (hAxis : Spec.Shape.nonemptyAxis? axis s = some hAxisPf) (hOut : TorchLean.Tensor.shapeAfterSum s axis = n.outShape) : g.evalAt payload input vals i = Except.ok { shape := n.outShape, tensor := hOut ▸ TorchLean.Tensor.reduceSum axis pT ⋯ }

evalAt_axisReduction_ok specialized to summation.

theorem Runtime.Autograd.IRExec.evalAt_reduceMean_ok {α : Type} [TorchLean.Storage α] [Context α] (g : NN.IR.Graph) (payload : NN.IR.Payload α) (input : Spec.SomeTensor α) (vals : Array (Spec.SomeTensor α)) (i : ℕ) (n : NN.IR.Node) (pId axis : ℕ) (s : Spec.Shape) (pT : TorchLean.Tensor α s) (hAxisPf : PLift (Spec.Shape.NonemptyAxis axis s)) (hN : g.getNode i = Except.ok n) (hk : n.kind = NN.IR.OpKind.reduceMean axis) (hp : NN.IR.unaryParent? n.parents = some pId) (hGet : vals[pId]? = some { shape := s, tensor := pT }) (hAxis : Spec.Shape.nonemptyAxis? axis s = some hAxisPf) (hOut : TorchLean.Tensor.shapeAfterSum s axis = n.outShape) : g.evalAt payload input vals i = Except.ok { shape := n.outShape, tensor := hOut ▸ TorchLean.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.IRExec.execOfState {α : Type} [TorchLean.Storage α] (inShape : Spec.Shape) (st : Internal.State α inShape) : ForwardGraph α

Repackage a lowered State as an ForwardGraph so we can call its evaluator helpers.

def Runtime.Autograd.IRExec.denoteAllState {α : Type} [TorchLean.Storage α] [Context α] (inShape : Spec.Shape) (st : Internal.State α inShape) (x : TorchLean.Tensor α inShape) : Array (Spec.SomeTensor α)

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

theorem Runtime.Autograd.IRExec.denoteAllState_snoc {α : Type} [TorchLean.Storage α] [Context α] {inShape : Spec.Shape} {ss : List Spec.Shape} {τ : Spec.Shape} (gd : ForwardData α [inShape] ss) (nodeData : ForwardNode α ([inShape] ++ ss) τ) (x : TorchLean.Tensor α inShape) : have st := ⟨ss, gd⟩; have st' := ⟨ss ++ [τ], gd.snoc nodeData⟩; denoteAllState inShape st' x = (denoteAllState inShape st x).push { shape := τ, tensor := nodeData.eval (gd.eval (TorchLean.TensorPack.cons x TorchLean.TensorPack.nil)) }

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

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

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

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

The forward executor's 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.IRExec.denoteAllState_get_mkIdx? {α : Type} [TorchLean.Storage α] [Context α] {inShape : Spec.Shape} {ss : List Spec.Shape} (gd : ForwardData α [inShape] ss) (x : TorchLean.Tensor α inShape) {pid : ℕ} {s : Spec.Shape} {idx : Proofs.Idx ([inShape] ++ ss) s} (hIdx : Internal.mkIdx inShape ss pid s = Except.ok idx) : (denoteAllState inShape ⟨ss, gd⟩ x)[pid]? = some { shape := s, tensor := Proofs.getIdx (gd.eval (TorchLean.TensorPack.cons x TorchLean.TensorPack.nil)) idx }

Lookup in denoteAllState agrees with getIdx when mkIdx pid s succeeds.

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

theorem Runtime.Autograd.IRExec.buildFrom_denoteAllFrom_finish {α : Type} [TorchLean.Storage α] [Context α] (g : NN.IR.Graph) (payload : NN.IR.Payload α) {inShape : Spec.Shape} {ss : List Spec.Shape} (i : ℕ) (x : TorchLean.Tensor α inShape) (hi : i < g.nodes.size) (τ : Spec.Shape) (nodeData : ForwardNode α ([inShape] ++ ss) τ) (st1 st' : Internal.State α inShape) (ctx : TorchLean.TensorPack α ([inShape] ++ ss)) (vals0 : Array (Spec.SomeTensor α)) (input : Spec.SomeTensor α) (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 { shape := τ, tensor := nodeData.eval ctx }) (hStep : denoteAllState inShape st1 x = vals0.push { shape := τ, tensor := nodeData.eval 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 forward-graph node’s forward (hEval), and
• the forward-graph table at i is the previous table plus the pushed node value (hStep), then denoteAllFrom at i returns the final forward-graph table.