Documentation

LeanMlir.Proofs.Foundation.IR

The small-net backward IR — `Back` / `Fwd` and their denotations #

A backward graph as an expression tree rooted at the cotangent (Back, Back3) and a forward graph (Fwd), each with a denotation ⟦·⟧ in the proofs' own Vec type, plus the bridges proving each emitted graph denotes the proven HasVJP.backward:

dense_back_bridge — dense's input gradient is one dot_general, so the bridge is definitional.
relu_back_bridge — compare(x > 0) + select; off the kink it denotes the canonical pdiv-derived ReLU backward (via relu_codegen_matches_canonical), which is exactly the codegen trust boundary.
the conv / maxpool / loss-cotangent bridges the Chapter 2–4 folds use (MlpTrainStep, CnnChainClose, CifarFold), and emitLossCot, which StableHLO's linear cotangent graph is proved against.

This is the IR the small nets' ties are stated in. ResNet onward states its ties about StableHLO's SHlo.

Design notes: the backward is modelled as an expression tree rooted at the cotangent — SSA/sharing is a semantics-preserving printer concern, so the correctness proof never touches it. The IR uses Vec/Mat directly rather than a 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 : ℕ) :

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

Instances For

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

Instances For

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: a dotGeneral node denotes post-composition with Mat.mulVec. Whole-network bridges chain lemmas of this shape, mirroring how vjpComp 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)

Instances For

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 vjpComp — 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

Instances For

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 = (denseHasVJP 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

Instances For

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 = (reluHasVJP 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) :

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

Equations

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

Instances For

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)

Instances For

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 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

Instances For

theorem Proofs.IR.reverseSlab_eq_gradSlab {h w kH kW : ℕ} (hkH : 2 * ((kH - 1) / 2) + 1 = kH) (hkW : 2 * ((kW - 1) / 2) + 1 = kW) (K : Fin kH → Fin kW → ℝ) (D : Fin h → Fin w → ℝ) (hi : Fin h) (wi : Fin w) :

(∑ kh : Fin kH, ∑ kw : Fin kW, K (kRev kh) (kRev kw) * if hpad : (kH - 1) / 2 ≤ ↑kh + ↑hi ∧ ↑kh + ↑hi - (kH - 1) / 2 < h ∧ (kW - 1) / 2 ≤ ↑kw + ↑wi ∧ ↑kw + ↑wi - (kW - 1) / 2 < w then D ⟨↑kh + ↑hi - (kH - 1) / 2, ⋯⟩ ⟨↑kw + ↑wi - (kW - 1) / 2, ⋯⟩ else 0) = ∑ ho : Fin h, ∑ wo : Fin w, if hpad : ↑ho ≤ ↑hi + (kH - 1) / 2 ∧ ↑hi + (kH - 1) / 2 - ↑ho < kH ∧ ↑wo ≤ ↑wi + (kW - 1) / 2 ∧ ↑wi + (kW - 1) / 2 - ↑wo < kW then K ⟨↑hi + (kH - 1) / 2 - ↑ho, ⋯⟩ ⟨↑wi + (kW - 1) / 2 - ↑wo, ⋯⟩ * D ho wo else 0

One channel slab of the reversed-kernel identity (odd kernels). For a kernel slab K and a cotangent slab D, correlating D with the spatially reversed K gives the input-gradient sum over output positions. Both sides range over the SAME 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 conv gate below is this at each channel pair, the depthwise one (depthwiseConv2d_dwReverse_eq_input_grad_formula) at each channel.

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 = conv2dInputGradFormula 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 conv2dInputGradFormula W, for ARBITRARY ic oc h w kH kW with odd kH, kW (2·⌊(kH-1)/2⌋+1 = kH). convFlatBack_eq_vjp_backward, which ties convFlatBack to the certified conv input-VJP, routes through this.

theorem Proofs.IR.conv_back_bridge {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) (b : Vec oc) (x : Tensor3 ic h w) (dy : Tensor3 oc h w) :

convBackDenote W dy = (conv2dHasVJP3 W b).backward x dy

Conv backward bridge (odd kernels, all dims). The emitted transposed-convolution graph denotes the proven conv input-VJP (conv2dHasVJP3 W b).backward: convBackDenote_eq_input_grad_formula read against the layer's VJP.

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 select_and_scatter, GE select): route each window's dy to its first maximal cell in row-major order (maxPool2Argmax), zero elsewhere. At a tie only that one cell receives dy, as in the printed op.

Equations

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

Instances For

theorem Proofs.IR.maxPool2Argmax_eq_iff_isArgmax {c h w : ℕ} (x : Tensor3 c (2 * h) (2 * w)) (h_smooth : MaxPool2Smooth x) (ci : Fin c) (hi_in : Fin (2 * h)) (wi_in : Fin (2 * w)) :

maxPool2Argmax x ci (winRow hi_in) (winCol wi_in) = (winRowMod hi_in, winColMod wi_in) ↔ MaxPool2IsArgmax x ci hi_in wi_in

Off a tie the first-argmax test is the argmax test. A cell the first argmax names is maximal at every point; under MaxPool2Smooth a maximal cell is the only one, hence the first.

theorem Proofs.IR.maxPoolBackDenote_eq_of_smooth {c h w : ℕ} (x : Tensor3 c (2 * h) (2 * w)) (h_smooth : MaxPool2Smooth x) (dy : Tensor3 c h w) :

maxPoolBackDenote x dy = (maxPool2HasVJPAt3 x h_smooth).backward dy

The first-argmax routing at a smooth point is maxPool2HasVJPAt3's backward.

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 = maxPool2HasVJP3.backward x dy ci hi_in wi_in

MaxPool backward bridge (smooth point). The emitted select_and_scatter 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. At a tie the denotation still routes to the printed op's one cell; only the VJP claim needs MaxPool2Smooth, since the pool has no derivative there.

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

Instances For

theorem Proofs.IR.gelu_back_bridge (n : ℕ) (x dy : Vec n) :

(emitActBack fun (i : Fin n) => geluScalarDeriv (x i)).denote dy = (geluHasVJP 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.geluErf_back_bridge (n : ℕ) (x dy : Vec n) :

(emitActBack fun (i : Fin n) => geluErfScalarDeriv (x i)).denote dy = (geluErfHasVJP n).backward x dy

Exact-GELU backward bridge. Same diagonal pattern, at the derivative of x · Φ(x).

theorem Proofs.IR.swish_back_bridge (n : ℕ) (x dy : Vec n) :

(emitActBack fun (i : Fin n) => swishScalarDeriv (x i)).denote dy = (swishHasVJP 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 = (sigmoidHasVJP 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.

Instances For

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 = (bnHasVJP n ε γ β hε).backward x dy

Full BatchNorm backward bridge. bnHasVJP = vjpComp normalize affine, so the emitted graph is the 3-term normalize graph fed γ ⊙ dy (the affine backward). Denotes (bnHasVJP …).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 = (layerNormHasVJP n ε γ β hε).backward x dy

LayerNorm backward bridge — free. layerNormHasVJP is definitionally bnHasVJP (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)

Instances For

theorem Proofs.IR.softmax_back_bridge (c : ℕ) (z dy : Vec c) :

(emitSoftmaxBack (softmax c z)).denote dy = (softmaxHasVJP 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.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 = (seBlockHasVJP 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 seBlockHasVJP.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)

Instances For

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)

Instances For

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)

Instances For

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 = (maxPool2HasVJPAt3 x h_smooth).backward dy

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

theorem Proofs.IR.conv3_node_bridge {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) (b : Vec oc) (x : Tensor3 ic h w) (dy : Tensor3 oc h w) :

(Back3.conv W Back3.cot).denote dy = (conv2dHasVJP3 W b).backward x dy

The Back3 conv node denotes the proven conv backward, for odd kH, kW (via conv_back_bridge).

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.

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

Instances For

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 = (maxPoolFlatHasVJPAt x h_smooth).backward dy

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

theorem Proofs.IR.conv_flatten_bridge {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) (b : Vec oc) (v : Vec (ic * h * w)) (dy : Vec (oc * h * w)) :

(Back3.conv W Back3.cot).flatDenote dy = (conv2dHasVJP3 W b).toHasVJP.backward v dy

Flatten bridge, conv (odd kernels). The flattened Back3 conv graph denotes the proven flattened conv layer backward HasVJP3.toHasVJP (conv2dHasVJP3 W b) — chains conv_back_bridge (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 = (reluHasVJPAt n x h_smooth).backward dy

ReLU _at bridge. The compare/select graph denotes the pointwise reluHasVJPAt 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 = ((denseHasVJP W b).toHasVJPAt v).backward dy

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

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.

Instances For

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 = (mlpHasVJPAt 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 mlpHasVJPAt.backward — the full assembly: per-op _at bridges chained through denote_subst, matching the nested vjpCompAt. 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)

Instances For

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

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

Instances For

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 denseWeightGrad. Certified by denseWeightGrad_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 denseBiasGrad. Certified by denseBiasGrad_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₂)

Instances For

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 (layer 1). 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 : ℕ) :

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

Instances For

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)

Instances For

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

Instances For

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) :

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

Instances For

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 vector emitLossCot equals ∂(crossEntropy)/∂logits (softmaxCE_grad), so the backward's cotangent leaf is proof-backed rather than supplied.