A denoted StableHLO-subset IR — Phase 0a/0b spike

Spike for planning/archive/typed_ir.md: give the emitted backward graph a denotational semantics ⟦·⟧ landing in the proofs' own Vec type, then prove the emitted graph denotes the proven HasVJP.backward. This turns the per-op proof↔codegen correspondence from a comment into a theorem.

This file is the scaffolding probe, not the full ladder:

• Phase 0a (dense) — dense_back_bridge. Dense's input-gradient is a single dot_general, so the bridge is definitional; its job is to pin the Back / denote plumbing.
• Phase 0b (relu, smooth point) — relu_back_bridge. The ReLU backward graph is compare(x > 0) + select; at a point off the kink it denotes the canonical pdiv-derived ReLU backward. This is the real content: it is conditional on smoothness and reuses the existing relu_codegen_matches_canonical, exactly matching the codegen trust boundary.

Design notes (see planning/archive/typed_ir.md): the backward is modelled as an expression tree rooted at the cotangent — SSA/sharing is a semantics-preserving printer concern (D2), so the correctness proof never touches it. The spike uses Vec/Mat directly (D1 shortcut) rather than the general flat-tensor type.

Everything closes under [propext, Classical.choice, Quot.sound] (audited in tests/AuditAxioms.lean); no native_decide.

inductive Proofs.IR.Back (inp : ℕ) : ℕ → Type

A backward subgraph, rooted at the cotangent dy : Vec inp, producing a Vec out. Saved forward data (weights A, the ReLU pre-activation x) is baked into the constructors. Each constructor models the StableHLO op a backward pass uses:

• cotangent — the graph input dy,
• dotGeneral — stablehlo.dot_general (here: matrix · vector),
• selectPos — stablehlo.compare GT 0 + stablehlo.select.

• cotangent {inp : ℕ} : Back inp inp
• dotGeneral {inp m n : ℕ} (A : Mat m n) : Back inp n → Back inp m
• selectPos {inp n : ℕ} (x : Vec n) : Back inp n → Back inp n
• scale {inp n : ℕ} (s : Vec n) : Back inp n → Back inp n
• sumBroadcast {inp n : ℕ} : Back inp n → Back inp n
• sub {inp n : ℕ} : Back inp n → Back inp n → Back inp n
• scaleConst {inp n : ℕ} (c : ℝ) : Back inp n → Back inp n
• add {inp n : ℕ} : Back inp n → Back inp n → Back inp n

noncomputable def Proofs.IR.Back.denote {inp out : ℕ} (e : Back inp out) (dy : Vec inp) : Vec out

Denotational semantics of a backward graph, into the proofs' own Vec type — so a bridge theorem can equate it with a proven HasVJP.backward.

Equations

• Proofs.IR.Back.cotangent.denote dy = dy
• (Proofs.IR.Back.dotGeneral A e').denote dy = A.mulVec (e'.denote dy)
• (Proofs.IR.Back.selectPos x e').denote dy = fun (i : Fin out) => if x i > 0 then e'.denote dy i else 0
• (Proofs.IR.Back.scale s e').denote dy = fun (i : Fin out) => e'.denote dy i * s i
• e'.sumBroadcast.denote dy = fun (x : Fin out) => ∑ j : Fin out, e'.denote dy j
• (e1.sub e2).denote dy = fun (i : Fin out) => e1.denote dy i - e2.denote dy i
• (Proofs.IR.Back.scaleConst c e').denote dy = fun (i : Fin out) => c * e'.denote dy i
• (e1.add e2).denote dy = fun (i : Fin out) => e1.denote dy i + e2.denote dy i

theorem Proofs.IR.denote_dotGeneral {inp m n : ℕ} (A : Mat m n) (e : Back inp n) (dy : Vec inp) : (Back.dotGeneral A e).denote dy = A.mulVec (e.denote dy)

Composition lemma (the Phase-3 mechanism in miniature): a dotGeneral node denotes post-composition with Mat.mulVec. Whole-network bridges will chain lemmas of this shape, mirroring how vjp_comp builds whole-net VJPs from per-layer ones.

def Proofs.IR.Back.subst {inp inp' out : ℕ} (e : Back inp out) (g : Back inp' inp) : Back inp' out

Plug the backward graph g into the cotangent leaf of e. For a composite g_layer ∘ f_layer, e is f_layer's backward and g is g_layer's, giving the composite's backward.

Equations

• Proofs.IR.Back.cotangent.subst g = g
• (Proofs.IR.Back.dotGeneral A e').subst g = Proofs.IR.Back.dotGeneral A (e'.subst g)
• (Proofs.IR.Back.selectPos x e').subst g = Proofs.IR.Back.selectPos x (e'.subst g)
• (Proofs.IR.Back.scale s e').subst g = Proofs.IR.Back.scale s (e'.subst g)
• e'.sumBroadcast.subst g = (e'.subst g).sumBroadcast
• (e1.sub e2).subst g = (e1.subst g).sub (e2.subst g)
• (Proofs.IR.Back.scaleConst c e').subst g = Proofs.IR.Back.scaleConst c (e'.subst g)
• (e1.add e2).subst g = (e1.subst g).add (e2.subst g)

theorem Proofs.IR.denote_subst {inp inp' out : ℕ} (e : Back inp out) (g : Back inp' inp) (dz : Vec inp') : (e.subst g).denote dz = e.denote (g.denote dz)

IR-level chain rule. subst denotes the composition of denotations: ⟦e[g/cotangent]⟧ dz = ⟦e⟧ (⟦g⟧ dz). The analogue of vjp_comp — chains per-op bridges into a whole-network bridge.

def Proofs.IR.emitDenseBack {m n : ℕ} (W : Mat m n) : Back n m

The dense input-gradient backward graph: one dot_general of the weight matrix with the cotangent.

Equations

• Proofs.IR.emitDenseBack W = Proofs.IR.Back.dotGeneral W Proofs.IR.Back.cotangent

theorem Proofs.IR.dense_back_bridge {m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (dy : Vec n) : (emitDenseBack W).denote dy = (dense_has_vjp W b).backward x dy

Dense bridge. The emitted graph denotes the proven dense backward Mat.mulVec W dy. Base case — dense's backward is a single dot_general — so this is definitional; it pins the plumbing.

def Proofs.IR.emitReluBack {n : ℕ} (x : Vec n) : Back n n

The ReLU backward graph: compare(x > 0) then select on the cotangent. x is the saved forward pre-activation.

Equations

• Proofs.IR.emitReluBack x = Proofs.IR.Back.selectPos x Proofs.IR.Back.cotangent

theorem Proofs.IR.relu_back_bridge {n : ℕ} (x : Vec n) (h_smooth : ∀ (k : Fin n), x k ≠ 0) (dy : Vec n) (i : Fin n) : (emitReluBack x).denote dy i = (relu_has_vjp n).backward x dy i

ReLU bridge (smooth point). At a point with no coordinate on the kink (∀ k, x k ≠ 0), the emitted compare/select graph denotes the canonical pdiv-derived ReLU backward. The real content: conditional on smoothness, reusing relu_codegen_matches_canonical. The Lean-vs-codegen gap at the kink is exactly the codegen trust boundary — and exactly where this equality is allowed to fail.

def Proofs.IR.kRev {k : ℕ} (i : Fin k) : Fin k

Spatial reversal of a kernel index: k − 1 − i.

Equations

• Proofs.IR.kRev i = ⟨k - 1 - ↑i, ⋯⟩

noncomputable def Proofs.IR.reverseSwap {ic oc kH kW : ℕ} (W : Kernel4 oc ic kH kW) : Kernel4 ic oc kH kW

Transpose-and-flip a kernel — swap in/out channels and reverse both spatial axes. This is the kernel the codegen feeds to the backward stablehlo.convolution (transpose dims [1,0,2,3] + reverse [2,3]).

Equations

• Proofs.IR.reverseSwap W ci co kh kw = W co ci (Proofs.IR.kRev kh) (Proofs.IR.kRev kw)

noncomputable def Proofs.IR.convBackDenote {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) : Tensor3 oc h w → Tensor3 ic h w

Denotation of the emitted conv input-gradient graph. The codegen emits convolution(dy, reverse(transpose(W))); under ⟦conv⟧ := conv2d (D3) that denotes a forward conv2d of the reversed-swapped kernel.

Equations

• Proofs.IR.convBackDenote W = Proofs.conv2d (Proofs.IR.reverseSwap W) fun (x : Fin ic) => 0

theorem Proofs.IR.convBackDenote_eq_input_grad_formula {ic oc h w kH kW : ℕ} (hkH : 2 * ((kH - 1) / 2) + 1 = kH) (hkW : 2 * ((kW - 1) / 2) + 1 = kW) (W : Kernel4 oc ic kH kW) (dy : Tensor3 oc h w) : conv2d (reverseSwap W) (fun (x : Fin ic) => 0) dy = conv2d_input_grad_formula W dy

The general conv-adjoint identity (odd kernels), all dims. The emitted reversed-kernel forward conv conv2d (reverseSwap W) 0 equals the certified conv input-gradient conv2d_input_grad_formula W, for ARBITRARY ic oc h w kH kW with odd kH, kW (2·⌊(kH-1)/2⌋+1 = kH). This is the reversed-kernel ⇒ correlation-adjoint reindex that conv_back_bridge_{1to2,2to2} previously asserted only at two toy 4×4 shapes by exhaustive fin_cases.

Proof: per output coordinate, both sides sum over the input channel co; the inner (kh,kw) sum (LHS, over the kernel window) and the (ho,wo) sum (RHS, over output positions) range over the SAME set of valid alignments via the partial bijection (kh,kw) ↦ (kh+hi-pH, kw+wi-pW) on the pad supports. Under oddness 2·pH = kH-1, the reversed-kernel index kH-1-kh matches the formula's hi+pH-ho, and the data indices coincide — so the matched summands are equal. Finset.sum_bij' over the pad-filtered supports; all index arithmetic by omega.

The single load-bearing leaf for the §B certified-VJP tie: every conv-heavy net's backward (convFlatBack) routes its conv input-grad through this.

theorem Proofs.IR.conv_back_bridge_1to2 (W : Kernel4 2 1 3 3) (b : Vec 2) (x : Tensor3 1 (2 * 2) (2 * 2)) (dy : Tensor3 2 (2 * 2) (2 * 2)) : convBackDenote W dy = (conv2d_has_vjp3 W b).backward x dy

Conv backward bridge, 1→2 channels (the Spatial instance's first conv: Kernel4 2 1 3 3 at 4×4). The emitted transposed-convolution graph denotes the proven conv input-VJP (conv2d_has_vjp3 W b).backward. Now a one-line instance of the general convBackDenote_eq_input_grad_formula (3×3 is odd) — no longer the brute-force fin_cases expansion.

theorem Proofs.IR.conv_back_bridge_2to2 (W : Kernel4 2 2 3 3) (b : Vec 2) (x dy : Tensor3 2 (2 * 2) (2 * 2)) : convBackDenote W dy = (conv2d_has_vjp3 W b).backward x dy

Conv backward bridge, 2→2 channels (the Spatial instance's second conv: Kernel4 2 2 3 3 at 4×4). Same identity at the 2→2 shape — also a one-line instance of the general lemma.

noncomputable def Proofs.IR.maxPoolBackDenote {c h w : ℕ} (x : Tensor3 c (2 * h) (2 * w)) : Tensor3 c h w → Tensor3 c (2 * h) (2 * w)

Denotation of the emitted maxpool input-gradient graph (StableHLO tile-compare-select): at a smooth point, route dy to each window's argmax input cell, zero elsewhere.

Equations

• Proofs.IR.maxPoolBackDenote x dy ci hi_in wi_in = if Proofs.MaxPool2IsArgmax x ci hi_in wi_in then dy ci (Proofs.winRow hi_in) (Proofs.winCol wi_in) else 0

theorem Proofs.IR.maxpool_back_bridge {c h w : ℕ} (x : Tensor3 c (2 * h) (2 * w)) (h_smooth : MaxPool2Smooth x) (dy : Tensor3 c h w) (ci : Fin c) (hi_in : Fin (2 * h)) (wi_in : Fin (2 * w)) : maxPoolBackDenote x dy ci hi_in wi_in = maxPool2_has_vjp3.backward x dy ci hi_in wi_in

MaxPool backward bridge (smooth point). The emitted tile-compare-select graph denotes the canonical pdiv-derived maxpool backward, conditional on no argmax ties (MaxPool2Smooth). The spatial-pooling analogue of relu_back_bridge; reuses maxPool2_codegen_matches_canonical. The Lean-vs-codegen gap at argmax-tie boundaries is exactly the codegen trust boundary.

def Proofs.IR.emitActBack {n : ℕ} (s : Vec n) : Back n n

The emitted elementwise-activation backward graph: stablehlo.multiply of the cotangent with the saved activation-derivative s = act'(x).

Equations

• Proofs.IR.emitActBack s = Proofs.IR.Back.scale s Proofs.IR.Back.cotangent

theorem Proofs.IR.gelu_back_bridge (n : ℕ) (x dy : Vec n) : (emitActBack fun (i : Fin n) => geluScalarDeriv (x i)).denote dy = (gelu_has_vjp n).backward x dy

GELU backward bridge. The emitted dy ⊙ gelu'(x) graph denotes the proven GELU backward. Definitional — GELU's diagonal Jacobian makes its HasVJP.backward exactly this elementwise scaling.

theorem Proofs.IR.swish_back_bridge (n : ℕ) (x dy : Vec n) : (emitActBack fun (i : Fin n) => swishScalarDeriv (x i)).denote dy = (swish_has_vjp n).backward x dy

Swish / SiLU backward bridge. Same diagonal pattern.

theorem Proofs.IR.sigmoid_back_bridge (n : ℕ) (x dy : Vec n) : (emitActBack fun (i : Fin n) => sigmoidScalarDeriv (x i)).denote dy = (sigmoid_has_vjp n).backward x dy

Sigmoid backward bridge. Same diagonal pattern.

noncomputable def Proofs.IR.bnNormalizeBackOf {n : ℕ} (xh : Vec n) (s invN : ℝ) (input : Back n n) : Back n n

The emitted BN-normalize input-gradient graph, as a function of an input subgraph (the cotangent for normalize alone; γ ⊙ dy for full BN). Two sumBroadcast reductions + sub + scaleConst — exactly the consolidated three-term formula.

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.IR.bn_affine_back_bridge {n : ℕ} (γ β : ℝ) (v dy : Vec n) : (Back.scaleConst γ Back.cotangent).denote dy = (bnAffine_has_vjp n γ β).backward v dy

BN affine backward bridge — the γ·dy half is one scaleConst.

theorem Proofs.IR.bn_normalize_back_bridge {n : ℕ} (ε : ℝ) (hε : 0 < ε) (x dxhat : Vec n) : (bnNormalizeBackOf (bnXhat n ε x) (bnIstd n x ε) (1 / ↑n) Back.cotangent).denote dxhat = (bnNormalize_has_vjp n ε hε).backward x dxhat

BN normalize backward bridge — the 3-term rank-1 wringer. The emitted reduce+broadcast+elementwise graph denotes the proven consolidated BN-normalize backward bnNormalize_has_vjp.backward (the cross-coordinate Σ dx̂ and Σ x̂·dx̂ reductions and the rank-1 x̂ᵢ·Σx̂·dx̂ correction, matched termwise).

theorem Proofs.IR.bn_back_bridge {n : ℕ} (ε γ β : ℝ) (hε : 0 < ε) (x dy : Vec n) : (bnNormalizeBackOf (bnXhat n ε x) (bnIstd n x ε) (1 / ↑n) (Back.scaleConst γ Back.cotangent)).denote dy = (bn_has_vjp n ε γ β hε).backward x dy

Full BatchNorm backward bridge. bn_has_vjp = vjp_comp normalize affine, so the emitted graph is the 3-term normalize graph fed γ ⊙ dy (the affine backward). Denotes (bn_has_vjp …).backward. The bnForward = bnAffine ∘ bnNormalize cast collapses by rfl.

theorem Proofs.IR.layernorm_back_bridge {n : ℕ} (ε γ β : ℝ) (hε : 0 < ε) (x dy : Vec n) : (bnNormalizeBackOf (bnXhat n ε x) (bnIstd n x ε) (1 / ↑n) (Back.scaleConst γ Back.cotangent)).denote dy = (layerNorm_has_vjp n ε γ β hε).backward x dy

LayerNorm backward bridge — free. layerNorm_has_vjp is definitionally bn_has_vjp (LayerNorm is BN on a different axis), so the same emitted graph denotes its backward.

noncomputable def Proofs.IR.emitSoftmaxBack {c : ℕ} (p : Vec c) : Back c c

The emitted softmax input-gradient graph: scale by p, subtract the broadcast inner product ⟨p, dy⟩, scale by p.

Equations

• Proofs.IR.emitSoftmaxBack p = Proofs.IR.Back.scale p (Proofs.IR.Back.cotangent.sub (Proofs.IR.Back.scale p Proofs.IR.Back.cotangent).sumBroadcast)

theorem Proofs.IR.softmax_back_bridge (c : ℕ) (z dy : Vec c) : (emitSoftmaxBack (softmax c z)).denote dy = (softmax_has_vjp c).backward z dy

Softmax backward bridge. The emitted reduce+broadcast+scale graph denotes the proven rank-1 softmax backward pᵢ·(dyᵢ − ⟨p, dy⟩).

theorem Proofs.IR.twoDense_back_bridge {d₀ d₁ d₂ : ℕ} (W₁ : Mat d₀ d₁) (b₁ : Vec d₁) (W₂ : Mat d₁ d₂) (b₂ : Vec d₂) (x : Vec d₀) (dz : Vec d₂) : ((emitDenseBack W₁).subst (emitDenseBack W₂)).denote dz = (vjp_comp (dense W₁ b₁) (dense W₂ b₂) ⋯ ⋯ (dense_has_vjp W₁ b₁) (dense_has_vjp W₂ b₂)).backward x dz

End-to-end composition bridge. The IR subst of two dense layers' backward graphs denotes the proven composite VJP (vjp_comp …).backward — denote_subst (IR chain rule) ∘ the per-op dense bridge.

theorem Proofs.IR.se_back_bridge {n : ℕ} (gate : Vec n → Vec n) (hg_diff : Differentiable ℝ gate) (hg : HasVJP gate) (bg : Back n n) (x dy : Vec n) (hbg : ∀ (z : Vec n), bg.denote z = hg.backward x z) : ((Back.scale (gate x) Back.cotangent).add (bg.subst (Back.scale x Back.cotangent))).denote dy = (seBlock_has_vjp gate hg_diff hg).backward x dy

Squeeze-and-Excitation backward bridge. Given the gate's backward graph bg bridged to its proven backward at x (hbg), the emitted SE backward graph — add (scale (gate x) dy) (bg[x ⊙ dy]) — denotes the proven seBlock_has_vjp.backward. The fan-in + denote_subst (to plug the gate graph in) assemble the per-op bridges through a non-composition combinator.

inductive Proofs.IR.Back3 (c₁ h₁ w₁ : ℕ) : ℕ → ℕ → ℕ → Type

Tensor3-level backward graph: indexed by the top cotangent shape (c₁ h₁ w₁) and the current shape (c₂ h₂ w₂) (walking toward the input gradient).

• cot {c₁ h₁ w₁ : ℕ} : Back3 c₁ h₁ w₁ c₁ h₁ w₁
• conv {c₁ h₁ w₁ ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) : Back3 c₁ h₁ w₁ oc h w → Back3 c₁ h₁ w₁ ic h w
• maxpool {c₁ h₁ w₁ c h w : ℕ} (x : Tensor3 c (2 * h) (2 * w)) : Back3 c₁ h₁ w₁ c h w → Back3 c₁ h₁ w₁ c (2 * h) (2 * w)

noncomputable def Proofs.IR.Back3.denote {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} (e : Back3 c₁ h₁ w₁ c₂ h₂ w₂) (dy : Tensor3 c₁ h₁ w₁) : Tensor3 c₂ h₂ w₂

Denote a Back3 graph as a Tensor3 → Tensor3 function, via the per-op Tensor3 backward denotations.

Equations

• Proofs.IR.Back3.cot.denote dy = dy
• (Proofs.IR.Back3.conv W e').denote dy = Proofs.IR.convBackDenote W (e'.denote dy)
• (Proofs.IR.Back3.maxpool x e').denote dy = Proofs.IR.maxPoolBackDenote x (e'.denote dy)

def Proofs.IR.Back3.subst {c₁ h₁ w₁ c₀ h₀ w₀ c₂ h₂ w₂ : ℕ} (e : Back3 c₁ h₁ w₁ c₂ h₂ w₂) (g : Back3 c₀ h₀ w₀ c₁ h₁ w₁) : Back3 c₀ h₀ w₀ c₂ h₂ w₂

Plug g into the cotangent leaf of e (Tensor3 composition).

Equations

• Proofs.IR.Back3.cot.subst g = g
• (Proofs.IR.Back3.conv W e').subst g = Proofs.IR.Back3.conv W (e'.subst g)
• (Proofs.IR.Back3.maxpool x e').subst g = Proofs.IR.Back3.maxpool x (e'.subst g)

theorem Proofs.IR.denote_subst3 {c₁ h₁ w₁ c₀ h₀ w₀ c₂ h₂ w₂ : ℕ} (e : Back3 c₁ h₁ w₁ c₂ h₂ w₂) (g : Back3 c₀ h₀ w₀ c₁ h₁ w₁) (dz : Tensor3 c₀ h₀ w₀) : (e.subst g).denote dz = e.denote (g.denote dz)

Tensor3 chain rule — the Back3 analogue of denote_subst.

theorem Proofs.IR.maxpool3_node_bridge {c h w : ℕ} (x : Tensor3 c (2 * h) (2 * w)) (h_smooth : MaxPool2Smooth x) (dy : Tensor3 c h w) : (Back3.maxpool x Back3.cot).denote dy = (maxPool2_has_vjp_at3 x h_smooth).backward dy

The Back3 maxpool node denotes the proven pointwise maxpool backward maxPool2_has_vjp_at3 — maxPoolBackDenote is that backward.

theorem Proofs.IR.conv3_node_bridge_1to2 (W : Kernel4 2 1 3 3) (b : Vec 2) (x : Tensor3 1 (2 * 2) (2 * 2)) (dy : Tensor3 2 (2 * 2) (2 * 2)) : (Back3.conv W Back3.cot).denote dy = (conv2d_has_vjp3 W b).backward x dy

The Back3 conv node denotes the proven conv backward, at the Spatial instance's 1→2 conv shape (via conv_back_bridge_1to2).

theorem Proofs.IR.conv_compose3 {ic mc oc h w kH₁ kW₁ kH₂ kW₂ : ℕ} (W₁ : Kernel4 mc ic kH₁ kW₁) (W₂ : Kernel4 oc mc kH₂ kW₂) (dz : Tensor3 oc h w) : ((Back3.conv W₁ Back3.cot).subst (Back3.conv W₂ Back3.cot)).denote dz = convBackDenote W₁ (convBackDenote W₂ dz)

Tensor3 composition demonstrator. The Back3 subst of two conv layers' backward graphs denotes the composition of their Tensor3 backwards, via the Tensor3 chain rule denote_subst3 — the conv/maxpool analogue of twoDense_back_bridge.

noncomputable def Proofs.IR.Back3.flatDenote {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} (e : Back3 c₁ h₁ w₁ c₂ h₂ w₂) (dy : Vec (c₁ * h₁ * w₁)) : Vec (c₂ * h₂ * w₂)

View a Back3 graph in flattened Vec space: flatten ∘ denote ∘ unflatten.

Equations

• e.flatDenote dy = (e.denote (Proofs.Tensor3.unflatten dy)).flatten

theorem Proofs.IR.maxpool_flatten_bridge {c h w : ℕ} (x : Tensor3 c (2 * h) (2 * w)) (h_smooth : MaxPool2Smooth x) (dy : Vec (c * h * w)) : (Back3.maxpool x Back3.cot).flatDenote dy = (maxPoolFlat_has_vjp_at x h_smooth).backward dy

Flatten bridge, max-pool. The flattened Back3 maxpool graph denotes the proven flattened maxpool layer backward maxPoolFlat_has_vjp_at (the form mnistCnnNoBn composes).

theorem Proofs.IR.conv_flatten_bridge_1to2 (W : Kernel4 2 1 3 3) (b : Vec 2) (v : Vec (1 * (2 * 2) * (2 * 2))) (dy : Vec (2 * (2 * 2) * (2 * 2))) : (Back3.conv W Back3.cot).flatDenote dy = (hasVJP3_to_hasVJP (conv2d_has_vjp3 W b)).backward v dy

Flatten bridge, conv (Spatial 1→2 shape). The flattened Back3 conv graph denotes the proven flattened conv layer backward hasVJP3_to_hasVJP (conv2d_has_vjp3 W b) — chains conv_back_bridge_1to2 (the reversed-kernel identity) with the Tensor3.flatten decode.

theorem Proofs.IR.relu_at_bridge (n : ℕ) (x : Vec n) (h_smooth : ∀ (k : Fin n), x k ≠ 0) (dy : Vec n) : (emitReluBack x).denote dy = (relu_has_vjp_at n x h_smooth).backward dy

ReLU _at bridge. The compare/select graph denotes the pointwise relu_has_vjp_at backward directly — definitional (no canonical sum).

theorem Proofs.IR.dense_at_bridge {m n : ℕ} (W : Mat m n) (b : Vec n) (v : Vec m) (dy : Vec n) : (emitDenseBack W).denote dy = ((dense_has_vjp W b).toHasVJPAt v).backward dy

Dense _at bridge. (dense_has_vjp).toHasVJPAt wraps the global instance, so the dense graph still denotes it (rfl).

theorem Proofs.IR.denseRelu_at_bridge {m n : ℕ} (W : Mat m n) (b : Vec n) (v : Vec m) (h_smooth : ∀ (k : Fin n), dense W b v k ≠ 0) (dy : Vec n) : ((emitDenseBack W).subst (emitReluBack (dense W b v))).denote dy = (vjp_comp_at (dense W b) (relu n) v ⋯ ⋯ ((dense_has_vjp W b).toHasVJPAt v) (relu_has_vjp_at n (dense W b v) h_smooth)).backward dy

Dense→ReLU block _at bridge. The IR subst of the dense and relu backward graphs denotes the proven vjp_comp_at block backward — a real mnistCnnNoBn building block, assembled from the per-op _at bridges via denote_subst.

noncomputable def Proofs.IR.emitMlpBack {d₀ d₁ d₂ d₃ : ℕ} (W₀ : Mat d₀ d₁) (W₁ : Mat d₁ d₂) (W₂ : Mat d₂ d₃) (p₀ : Vec d₁) (p₁ : Vec d₂) : Back d₃ d₀

The emitted backward graph for the whole MLP: the subst chain dense₀ ∘ relu(p₀) ∘ dense₁ ∘ relu(p₁) ∘ dense₂ (backward order), where p₀ = dense W₀ b₀ x, p₁ = dense W₁ b₁ (relu (dense W₀ b₀ x)) are the ReLU pre-activations.

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.IR.mlp_whole_bridge {d₀ d₁ d₂ d₃ : ℕ} (W₀ : Mat d₀ d₁) (b₀ : Vec d₁) (W₁ : Mat d₁ d₂) (b₁ : Vec d₂) (W₂ : Mat d₂ d₃) (b₂ : Vec d₃) (x : Vec d₀) (h_smooth_0 : ∀ (k : Fin d₁), dense W₀ b₀ x k ≠ 0) (h_smooth_1 : ∀ (k : Fin d₂), dense W₁ b₁ (relu d₁ (dense W₀ b₀ x)) k ≠ 0) (dy : Vec d₃) : (emitMlpBack W₀ W₁ W₂ (dense W₀ b₀ x) (dense W₁ b₁ (relu d₁ (dense W₀ b₀ x)))).denote dy = (mlp_has_vjp_at W₀ b₀ W₁ b₁ W₂ b₂ x h_smooth_0 h_smooth_1).backward dy

Whole-network bridge. The emitted MLP backward graph denotes the proven mlp_has_vjp_at.backward — the full assembly: per-op _at bridges chained through denote_subst, matching the nested vjp_comp_at. A machine-checked statement that the emitted backward graph computes the proven whole-network VJP at a smooth point.

noncomputable def Proofs.IR.emitWeightGrad {inp m n : ℕ} (x : Vec m) (e : Back inp n) (dy : Vec inp) : Mat m n

Weight-gradient emitter: the outer product of the dense layer's saved forward input x with the cotangent at the layer's output (the denotation of the backward subgraph e). Mirrors the dot_general that contracts the batch axis (dW = xᵀ · dy).

Equations

• Proofs.IR.emitWeightGrad x e dy = Proofs.Mat.outer x (e.denote dy)

noncomputable def Proofs.IR.emitBiasGrad {inp n : ℕ} (e : Back inp n) (dy : Vec inp) : Vec n

Bias-gradient emitter: the cotangent at the layer's output. Mirrors the reduce-add over the batch axis (db = Σ_batch dy).

Equations

• Proofs.IR.emitBiasGrad e dy = e.denote dy

theorem Proofs.IR.weight_grad_bridge {inp m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (e : Back inp n) (dy : Vec inp) (i : Fin m) (j : Fin n) : emitWeightGrad x e dy i j = ∑ k : Fin n, pdiv (fun (v : Vec (m * n)) => dense (Mat.unflatten v) b x) W.flatten (finProdFinEquiv (i, j)) k * e.denote dy k

Weight-gradient bridge. The emitted outer-product graph, fed the cotangent the backward subgraph e delivers, computes coordinate-wise the cotangent-contracted Jacobian of the dense layer wrt W — the proven dense_weight_grad. Certified by dense_weight_grad_correct at the actual chain cotangent e.denote dy, so it composes with any of the input-gradient bridges above.

theorem Proofs.IR.bias_grad_bridge {inp m n : ℕ} (W : Mat m n) (b : Vec n) (x : Vec m) (e : Back inp n) (dy : Vec inp) (i : Fin n) : emitBiasGrad e dy i = ∑ j : Fin n, pdiv (fun (b' : Vec n) => dense W b' x) b i j * e.denote dy j

Bias-gradient bridge. The emitted graph (the cotangent itself, reduce-summed over the batch) computes the cotangent-contracted Jacobian of the dense layer wrt b — the proven dense_bias_grad. Certified by dense_bias_grad_correct.

def Proofs.IR.mlpCotOut1 {d₂ d₃ : ℕ} (W₂ : Mat d₂ d₃) (p₁ : Vec d₂) : Back d₃ d₂

The backward subgraph delivering the cotangent at the MLP's layer-1 dense output: relu'(p₁) ⊙ (W₂ · dy) — ReLU-back composed with the layer-2 dense input-gradient. (Layer 2's output cotangent is the top cotangent; layer 0's prepends another relu'(p₀) ⊙ (W₁ · ·).)

Equations

• Proofs.IR.mlpCotOut1 W₂ p₁ = (Proofs.IR.emitReluBack p₁).subst (Proofs.IR.emitDenseBack W₂)

theorem Proofs.IR.mlp_layer1_weight_grad_bridge {d₁ d₂ d₃ : ℕ} (W₁ : Mat d₁ d₂) (b₁ : Vec d₂) (W₂ : Mat d₂ d₃) (x₁ : Vec d₁) (p₁ : Vec d₂) (dy : Vec d₃) (i : Fin d₁) (j : Fin d₂) : emitWeightGrad x₁ (mlpCotOut1 W₂ p₁) dy i j = ∑ k : Fin d₂, pdiv (fun (v : Vec (d₁ * d₂)) => dense (Mat.unflatten v) b₁ x₁) W₁.flatten (finProdFinEquiv (i, j)) k * (mlpCotOut1 W₂ p₁).denote dy k

MLP hidden-layer parameter-gradient bridge (representative). At the interesting layer — layer 1, whose output cotangent is a genuine backward subgraph mlpCotOut1, not just the top cotangent — the emitted weight and bias gradients equal the certified Jacobians of that dense layer wrt W₁/b₁, contracted with the cotangent the backward chain actually delivers there. Instantiates the generic bridges at the MLP's layer-1 subgraph; the other two layers are the same bridges at Back.cotangent (layer 2) and mlpCotOut0 (layer 0).

inductive Proofs.IR.Fwd (inp : ℕ) : ℕ → Type

A forward graph: input x : Vec inp, producing a Vec out. Each constructor is a forward op (the affine dense, the relu nonlinearity); the StableHLO mirror renders dense as dot_general + broadcast_in_dim + add and relu as maximum 0.

• input {inp : ℕ} : Fwd inp inp
• dense {inp m n : ℕ} (W : Mat m n) (b : Vec n) : Fwd inp m → Fwd inp n
• relu {inp n : ℕ} : Fwd inp n → Fwd inp n

noncomputable def Proofs.IR.Fwd.denote {inp out : ℕ} (e : Fwd inp out) (x : Vec inp) : Vec out

Denotational semantics of a forward graph, into the proofs' Vec type — so a bridge can equate it with the proven forward map.

Equations

• Proofs.IR.Fwd.input.denote x = x
• (Proofs.IR.Fwd.dense W b e').denote x = Proofs.dense W b (e'.denote x)
• e'.relu.denote x = Proofs.relu out (e'.denote x)

def Proofs.IR.Fwd.subst {inp mid out : ℕ} (e : Fwd mid out) (g : Fwd inp mid) : Fwd inp out

Plug g into the input leaf of e (forward composition). The forward analogue of Back.subst.

Equations

• Proofs.IR.Fwd.input.subst g = g
• (Proofs.IR.Fwd.dense W b e').subst g = Proofs.IR.Fwd.dense W b (e'.subst g)
• e'.relu.subst g = (e'.subst g).relu

theorem Proofs.IR.denote_subst_fwd {inp mid out : ℕ} (e : Fwd mid out) (g : Fwd inp mid) (x : Vec inp) : (e.subst g).denote x = e.denote (g.denote x)

Forward IR chain rule — subst denotes composition. Mirror of denote_subst; lets forward graphs compose to arbitrary depth.

def Proofs.IR.emitMlpFwd {d₀ d₁ d₂ d₃ : ℕ} (W₀ : Mat d₀ d₁) (b₀ : Vec d₁) (W₁ : Mat d₁ d₂) (b₁ : Vec d₂) (W₂ : Mat d₂ d₃) (b₂ : Vec d₃) : Fwd d₀ d₃

The emitted forward graph for the whole MLP: dense W₂ ∘ relu ∘ dense W₁ ∘ relu ∘ dense W₀ as a Fwd tree.

Equations

• Proofs.IR.emitMlpFwd W₀ b₀ W₁ b₁ W₂ b₂ = Proofs.IR.Fwd.dense W₂ b₂ (Proofs.IR.Fwd.dense W₁ b₁ (Proofs.IR.Fwd.dense W₀ b₀ Proofs.IR.Fwd.input).relu).relu

theorem Proofs.IR.mlp_fwd_bridge {d₀ d₁ d₂ d₃ : ℕ} (W₀ : Mat d₀ d₁) (b₀ : Vec d₁) (W₁ : Mat d₁ d₂) (b₁ : Vec d₂) (W₂ : Mat d₂ d₃) (b₂ : Vec d₃) (x : Vec d₀) : (emitMlpFwd W₀ b₀ W₁ b₁ W₂ b₂).denote x = mlpForward W₀ b₀ W₁ b₁ W₂ b₂ x

Forward bridge. The emitted forward graph denotes the proven forward map mlpForward. Promotes the train step's forward from trusted to proof-backed (up to the printer): the emitted forward StableHLO is the rendering of an IR proven to compute mlpForward.

theorem Proofs.IR.mlp_fwd_preact0 {d₀ d₁ : ℕ} (W₀ : Mat d₀ d₁) (b₀ : Vec d₁) (x : Vec d₀) : (Fwd.dense W₀ b₀ Fwd.input).denote x = dense W₀ b₀ x

Splice contract, layer-0 pre-activation. The forward IR's sub-graph up to the first ReLU input denotes exactly the pre-activation dense W₀ b₀ x that the backward graph reads as p₀ (its first compare/select mask). Forward output ↦ backward input, proven.

theorem Proofs.IR.mlp_fwd_preact1 {d₀ d₁ d₂ : ℕ} (W₀ : Mat d₀ d₁) (b₀ : Vec d₁) (W₁ : Mat d₁ d₂) (b₁ : Vec d₂) (x : Vec d₀) : (Fwd.dense W₁ b₁ (Fwd.dense W₀ b₀ Fwd.input).relu).denote x = dense W₁ b₁ (relu d₁ (dense W₀ b₀ x))

Splice contract, layer-1 pre-activation. Likewise the sub-graph up to the second ReLU input denotes the p₁ the backward reads.

noncomputable def Proofs.IR.emitLossCot (c : ℕ) (logits : Vec c) (label : Fin c) : Vec c

The emitted loss-cotangent (softmax-CE head): softmax(logits) − onehot(label). Rendered as exp + reduce(add) + broadcast + divide (softmax) then subtract the target. Feeds the backward's cotangent leaf.

Equations

• Proofs.IR.emitLossCot c logits label j = Proofs.softmax c logits j - Proofs.oneHot c label j

theorem Proofs.IR.lossCot_bridge (c : ℕ) (logits : Vec c) (label j : Fin c) : emitLossCot c logits label j = pdiv (fun (z : Vec c) (x : Fin 1) => crossEntropy c z label) logits j 0

Loss-cotangent bridge. The emitted softmax−onehot graph denotes the proven cross-entropy gradient ∂(crossEntropy)/∂logits (softmaxCE_grad). So the cotangent fed to the backward is itself proof-backed, not supplied: the whole train step forward → loss → backward → grads is proof-backed end to end, and only the SGD arithmetic (and printer/IREE/float) stays trusted.