The chapter 3–4 graphs denote the chapter nets #
StableHLO defines the MNIST-CNN and CIFAR forward graphs (cnnFwdGraph, cifarFwdGraph,
cifar8FwdGraph, cifar8BnFwdGraph) and the MNIST-CNN backward graph (cnnBackGraph) as SHlo
terms, and writes their artifacts. This file proves each denotes its net: the forwards are
mnistCnnNoBnForward / cifarCnnForward / cifarCnn8Forward / cifarCnnBn8Forward, and the
backward is the whole-network VJP mnistCnnNoBn_has_vjp_at.backward at a smooth point. Kept out of
StableHLO so the IR imports no net.
theorem
Proofs.StableHLO.cnnFwdGraph_faithful
{ic c h w d1 nClasses kH kW : ℕ}
(W₁ : Kernel4 c ic kH kW)
(b₁ : Vec c)
(W₂ : Kernel4 c c kH kW)
(b₂ : Vec c)
(W₃ : Mat (c * h * w) d1)
(b₃ : Vec d1)
(W₄ : Mat d1 d1)
(b₄ : Vec d1)
(W₅ : Mat d1 nClasses)
(b₅ : Vec nClasses)
(x : Vec (ic * (2 * h) * (2 * w)))
:
den (cnnFwdGraph W₁ b₁ W₂ b₂ W₃ b₃ W₄ b₄ W₅ b₅ x) = mnistCnnNoBnForward W₁ b₁ W₂ b₂ W₃ b₃ W₄ b₄ W₅ b₅ x
CNN forward faithfulness. The forward graph denotes the proven
mnistCnnNoBnForward.
theorem
Proofs.StableHLO.cifarFwdGraph_faithful
{ic c1 c2 h w d1 nClasses kH kW : ℕ}
(W₁ : Kernel4 c1 ic kH kW)
(b₁ : Vec c1)
(W₂ : Kernel4 c1 c1 kH kW)
(b₂ : Vec c1)
(W₃ : Kernel4 c2 c1 kH kW)
(b₃ : Vec c2)
(W₄ : Kernel4 c2 c2 kH kW)
(b₄ : Vec c2)
(W₅ : Mat (c2 * h * w) d1)
(b₅ : Vec d1)
(W₆ : Mat d1 d1)
(b₆ : Vec d1)
(W₇ : Mat d1 nClasses)
(b₇ : Vec nClasses)
(x : Vec (ic * (2 * (2 * h)) * (2 * (2 * w))))
:
den (cifarFwdGraph W₁ b₁ W₂ b₂ W₃ b₃ W₄ b₄ W₅ b₅ W₆ b₆ W₇ b₇ x) = cifarCnnForward W₁ b₁ W₂ b₂ W₃ b₃ W₄ b₄ W₅ b₅ W₆ b₆ W₇ b₇ x
CIFAR-CNN forward faithfulness. The forward graph denotes the proven
cifarCnnForward.
theorem
Proofs.StableHLO.cifar8FwdGraph_faithful
{ic c1 c2 c3 c4 h w d1 nClasses kH kW : ℕ}
(W₁ : Kernel4 c1 ic kH kW)
(b₁ : Vec c1)
(W₂ : Kernel4 c1 c1 kH kW)
(b₂ : Vec c1)
(W₃ : Kernel4 c2 c1 kH kW)
(b₃ : Vec c2)
(W₄ : Kernel4 c2 c2 kH kW)
(b₄ : Vec c2)
(W₅ : Kernel4 c3 c2 kH kW)
(b₅ : Vec c3)
(W₆ : Kernel4 c3 c3 kH kW)
(b₆ : Vec c3)
(W₇ : Kernel4 c4 c3 kH kW)
(b₇ : Vec c4)
(W₈ : Kernel4 c4 c4 kH kW)
(b₈ : Vec c4)
(W₉ : Mat (c4 * h * w) d1)
(b₉ : Vec d1)
(Wa : Mat d1 d1)
(ba : Vec d1)
(Wb : Mat d1 nClasses)
(bb : Vec nClasses)
(x : Vec (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))
:
den (cifar8FwdGraph W₁ b₁ W₂ b₂ W₃ b₃ W₄ b₄ W₅ b₅ W₆ b₆ W₇ b₇ W₈ b₈ W₉ b₉ Wa ba Wb bb x) = cifarCnn8Forward W₁ b₁ W₂ b₂ W₃ b₃ W₄ b₄ W₅ b₅ W₆ b₆ W₇ b₇ W₈ b₈ W₉ b₉ Wa ba Wb bb x
Deeper (8-conv) CIFAR-CNN forward faithfulness. The forward graph denotes the
proven cifarCnn8Forward.
theorem
Proofs.StableHLO.cifar8BnFwdGraph_faithful
{ic c1 c2 c3 c4 h w d1 nClasses kH kW : ℕ}
(epsStr : String)
(W₁ : Kernel4 c1 ic kH kW)
(b₁ : Vec c1)
(ε₁ : ℝ)
(γ₁ β₁ : Vec c1)
(W₂ : Kernel4 c1 c1 kH kW)
(b₂ : Vec c1)
(ε₂ : ℝ)
(γ₂ β₂ : Vec c1)
(W₃ : Kernel4 c2 c1 kH kW)
(b₃ : Vec c2)
(ε₃ : ℝ)
(γ₃ β₃ : Vec c2)
(W₄ : Kernel4 c2 c2 kH kW)
(b₄ : Vec c2)
(ε₄ : ℝ)
(γ₄ β₄ : Vec c2)
(W₅ : Kernel4 c3 c2 kH kW)
(b₅ : Vec c3)
(ε₅ : ℝ)
(γ₅ β₅ : Vec c3)
(W₆ : Kernel4 c3 c3 kH kW)
(b₆ : Vec c3)
(ε₆ : ℝ)
(γ₆ β₆ : Vec c3)
(W₇ : Kernel4 c4 c3 kH kW)
(b₇ : Vec c4)
(ε₇ : ℝ)
(γ₇ β₇ : Vec c4)
(W₈ : Kernel4 c4 c4 kH kW)
(b₈ : Vec c4)
(ε₈ : ℝ)
(γ₈ β₈ : Vec c4)
(W₉ : Mat (c4 * h * w) d1)
(b₉ : Vec d1)
(Wa : Mat d1 d1)
(ba : Vec d1)
(Wb : Mat d1 nClasses)
(bb : Vec nClasses)
(x : Vec (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))
:
den
(cifar8BnFwdGraph epsStr W₁ b₁ ε₁ γ₁ β₁ W₂ b₂ ε₂ γ₂ β₂ W₃ b₃ ε₃ γ₃ β₃ W₄ b₄ ε₄ γ₄ β₄ W₅ b₅ ε₅ γ₅ β₅ W₆ b₆ ε₆ γ₆ β₆
W₇ b₇ ε₇ γ₇ β₇ W₈ b₈ ε₈ γ₈ β₈ W₉ b₉ Wa ba Wb bb x) = cifarCnnBn8Forward W₁ b₁ ε₁ γ₁ β₁ W₂ b₂ ε₂ γ₂ β₂ W₃ b₃ ε₃ γ₃ β₃ W₄ b₄ ε₄ γ₄ β₄ W₅ b₅ ε₅ γ₅ β₅ W₆ b₆ ε₆ γ₆ β₆ W₇ b₇ ε₇
γ₇ β₇ W₈ b₈ ε₈ γ₈ β₈ W₉ b₉ Wa ba Wb bb x
Deeper (8-conv) BN-CIFAR forward faithfulness. The forward graph denotes the
proven cifarCnnBn8Forward.
noncomputable def
Proofs.StableHLO.maxPoolFlat_has_vjp_at'
{c h w : ℕ}
(v : Vec (c * (2 * h) * (2 * w)))
(hs : MaxPool2Smooth (Tensor3.unflatten v))
:
HasVJPAt (maxPoolFlat c h w) v
Max-pool VJP at a raw flattened point (no flatten ∘ unflatten index), so
it composes without a transport cast; backward is maxPoolBackFlat. The
correct field reuses maxPoolFlat_has_vjp_at.correct, aligning the point
via Tensor3.flatten_unflatten.
Equations
- Proofs.StableHLO.maxPoolFlat_has_vjp_at' v hs = { backward := Proofs.StableHLO.maxPoolBackFlat c h w v, correct := ⋯ }
Instances For
@[simp]
theorem
Proofs.StableHLO.maxPoolFlat_has_vjp_at'_backward
{c h w : ℕ}
(v : Vec (c * (2 * h) * (2 * w)))
(hs : MaxPool2Smooth (Tensor3.unflatten v))
:
theorem
Proofs.StableHLO.cnnBackGraph_faithful
{ic c h w d1 nClasses kH kW : ℕ}
(W₁ : Kernel4 c ic kH kW)
(b₁ : Vec c)
(W₂ : Kernel4 c c kH kW)
(b₂ : Vec c)
(W₃ : Mat (c * h * w) d1)
(b₃ : Vec d1)
(W₄ : Mat d1 d1)
(b₄ : Vec d1)
(W₅ : Mat d1 nClasses)
(b₅ : Vec nClasses)
(hc : 0 < c)
(hh : 0 < h)
(hw : 0 < w)
(x : Vec (ic * (2 * h) * (2 * w)))
(h1 : ∀ (k : Fin (c * (2 * h) * (2 * w))), flatConv W₁ b₁ x k ≠ 0)
(h2 : ∀ (k : Fin (c * (2 * h) * (2 * w))), flatConv W₂ b₂ ((relu (c * (2 * h) * (2 * w)) ∘ flatConv W₁ b₁) x) k ≠ 0)
(h_mp :
MaxPool2Smooth
(Tensor3.unflatten
(((relu (c * (2 * h) * (2 * w)) ∘ flatConv W₂ b₂) ∘ relu (c * (2 * h) * (2 * w)) ∘ flatConv W₁ b₁) x)))
(h3 :
∀ (k : Fin d1),
dense W₃ b₃
(maxPoolFlat c h w
(((relu (c * (2 * h) * (2 * w)) ∘ flatConv W₂ b₂) ∘ relu (c * (2 * h) * (2 * w)) ∘ flatConv W₁ b₁) x))
k ≠ 0)
(h4 :
∀ (k : Fin d1),
dense W₄ b₄
((relu d1 ∘ dense W₃ b₃)
(maxPoolFlat c h w
(((relu (c * (2 * h) * (2 * w)) ∘ flatConv W₂ b₂) ∘ relu (c * (2 * h) * (2 * w)) ∘ flatConv W₁ b₁) x)))
k ≠ 0)
(dy : Vec nClasses)
:
den (cnnBackGraph W₁ b₁ W₂ b₂ W₃ b₃ W₄ b₄ W₅ x dy) = (mnistCnnNoBn_has_vjp_at W₁ b₁ W₂ b₂ W₃ b₃ W₄ b₄ W₅ b₅ hc hh hw x h1 h2 h_mp h3 h4).backward dy