Documentation

LeanMlir.Proofs.Foundation.BatchedStages

Batched stages — conv/depthwise → true batch-norm (→ swish) at the batched index, with VJPs #

The building blocks every batched conv-net forward is written in, at the flat batched index N·(c·h·w): each stage is batchMap N of a proven per-example op, then true batch-norm (bnBatchLA, which couples the batch), then pointwise swish where the net has it.

stage	forward	VJP
1×1 / stride-1 conv → BN → swish	`cbsB`	`cbsB_has_vjp`
XLA-`SAME` stride-2 stem conv → BN → swish	`stemB`	`stemB_has_vjp`
depthwise (stride 1 / symmetric stride 2) → BN → swish	`dwbsB` / `dwbsSB`	`dwbsB_has_vjp` / `dwbsSB_has_vjp`
squeeze-excite	`seB`	`seB_has_vjp`
1×1 projection → BN (no activation)	`projB`	`projB_has_vjp`

The generic pieces they compose — batchMap_has_vjp (block-diagonal VJP of a batch-separable op) and bnBatchLA_has_vjp — are in BatchMapVJPAt. EfficientNet-B0, MobileNetV2/V4 and the ResNets build their blocks from these; the MBConv blocks themselves are in EfficientNetRenderPC and EfficientNetChainClose.

@[reducible]

noncomputable def Proofs.cbsB (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (γ β : Vec oc) :

Vec (N * (ic * h * w)) → Vec (N * (oc * h * w))

Batched conv → bn → swish (1×1 expand / generic stride-1 conv).

Equations

Proofs.cbsB N W b ε γ β = Proofs.swish (N * (oc * h * w)) ∘ Proofs.StableHLO.bnBatchLA N oc h w ε γ β ∘ Proofs.StableHLO.batchMap N (Proofs.flatConv W b)

Instances For

noncomputable def Proofs.stemB (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (γ β : Vec oc) :

Vec (N * (ic * (2 * h) * (2 * w))) → Vec (N * (oc * h * w))

Batched strided (3×3 s2) stem conv → bn → swish (halves spatial). ⚠ At the XLA-SAME phase (flatConvStride2Xla = decimateOddFlat ∘ flatConv): the TF-origin B0 pads its stem (0,1), and the shipped render has emitted convStridedXla there since 2026-08-08. The symmetric flatConvStride2 has the same type and output shape; nothing structural would notice the wrong one (re-spelled 2026-09-05, planning/archive/xla_same_respell_and_blueprint_audit.md).

Equations

Proofs.stemB N W b ε γ β = Proofs.swish (N * (oc * h * w)) ∘ Proofs.StableHLO.bnBatchLA N oc h w ε γ β ∘ Proofs.StableHLO.batchMap N (Proofs.flatConvStride2Xla W b)

Instances For

@[reducible]

noncomputable def Proofs.dwbsB (N : ℕ) {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ℝ) (γ β : Vec c) :

Vec (N * (c * h * w)) → Vec (N * (c * h * w))

Batched depthwise (stride-1, k×k) → bn → swish.

Equations

Proofs.dwbsB N W b ε γ β = Proofs.swish (N * (c * h * w)) ∘ Proofs.StableHLO.bnBatchLA N c h w ε γ β ∘ Proofs.StableHLO.batchMap N (Proofs.depthwiseFlat W b)

Instances For

@[reducible]

noncomputable def Proofs.dwbsSB (N : ℕ) {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ℝ) (γ β : Vec c) :

Vec (N * (c * (2 * h) * (2 * w))) → Vec (N * (c * h * w))

Batched depthwise (stride-2 downsample, k×k) → bn → swish.

Equations

Proofs.dwbsSB N W b ε γ β = Proofs.swish (N * (c * h * w)) ∘ Proofs.StableHLO.bnBatchLA N c h w ε γ β ∘ Proofs.StableHLO.batchMap N (Proofs.depthwiseStride2Flat W b)

Instances For

@[reducible]

noncomputable def Proofs.seB (N : ℕ) {c h w r : ℕ} (W₁ : Mat c r) (b₁ : Vec r) (W₂ : Mat r c) (b₂ : Vec c) :

Vec (N * (c * h * w)) → Vec (N * (c * h * w))

Batched squeeze-excite block x ⊙ gate(x) (the proven seBlockFull, per example).

Equations

Proofs.seB N W₁ b₁ W₂ b₂ = Proofs.StableHLO.batchMap N (Proofs.seBlockFull W₁ b₁ W₂ b₂)

Instances For

@[reducible]

noncomputable def Proofs.projB (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (γ β : Vec oc) :

Vec (N * (ic * h * w)) → Vec (N * (oc * h * w))

Batched project: 1×1 conv → bn (no swish — the linear bottleneck).

Equations

Proofs.projB N W b ε γ β = Proofs.StableHLO.bnBatchLA N oc h w ε γ β ∘ Proofs.StableHLO.batchMap N (Proofs.flatConv W b)

Instances For

noncomputable def Proofs.flatConv_has_vjp {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) :

HasVJP (flatConv W b)

flatConv W b VJP — the per-example 1×1/3×3 conv input-VJP (the HasVJP3-bridged conv2d).

Equations

Proofs.flatConv_has_vjp W b = Proofs.hasVJP3_to_hasVJP (Proofs.conv2d_has_vjp3 W b)

Instances For

theorem Proofs.bnSwishStage_differentiable (N : ℕ) {a oc h w : ℕ} (op : Vec a → Vec (oc * h * w)) (hop : Differentiable ℝ op) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

Differentiable ℝ (swish (N * (oc * h * w)) ∘ StableHLO.bnBatchLA N oc h w ε γ β ∘ StableHLO.batchMap N op)

Differentiability of a batched conv/depthwise → bn → swish stage.

noncomputable def Proofs.bnSwishStage_has_vjp (N : ℕ) {a oc h w : ℕ} (op : Vec a → Vec (oc * h * w)) (hop : Differentiable ℝ op) (hopv : HasVJP op) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

HasVJP (swish (N * (oc * h * w)) ∘ StableHLO.bnBatchLA N oc h w ε γ β ∘ StableHLO.batchMap N op)

VJP of a batched conv/depthwise → bn → swish stage: lift op's VJP per example, then the proven true-BN VJP, then the pointwise swish VJP.

Equations

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

Instances For

theorem Proofs.bnStage_differentiable (N : ℕ) {a oc h w : ℕ} (op : Vec a → Vec (oc * h * w)) (hop : Differentiable ℝ op) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

Differentiable ℝ (StableHLO.bnBatchLA N oc h w ε γ β ∘ StableHLO.batchMap N op)

Differentiability of a batched conv → bn stage (project bottleneck, no swish).

noncomputable def Proofs.bnStage_has_vjp (N : ℕ) {a oc h w : ℕ} (op : Vec a → Vec (oc * h * w)) (hop : Differentiable ℝ op) (hopv : HasVJP op) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

HasVJP (StableHLO.bnBatchLA N oc h w ε γ β ∘ StableHLO.batchMap N op)

VJP of a batched conv → bn stage.

Equations

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

Instances For

theorem Proofs.cbsB_differentiable (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

Differentiable ℝ (cbsB N W b ε γ β)

noncomputable def Proofs.cbsB_has_vjp (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

HasVJP (cbsB N W b ε γ β)

Equations

Proofs.cbsB_has_vjp N W b ε hε γ β = Proofs.bnSwishStage_has_vjp N (Proofs.flatConv W b) ⋯ (Proofs.flatConv_has_vjp W b) ε hε γ β

Instances For

theorem Proofs.stemB_differentiable (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

Differentiable ℝ (stemB N W b ε γ β)

noncomputable def Proofs.stemB_has_vjp (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

HasVJP (stemB N W b ε γ β)

Equations

Proofs.stemB_has_vjp N W b ε hε γ β = Proofs.bnSwishStage_has_vjp N (Proofs.flatConvStride2Xla W b) ⋯ (Proofs.flatConvStride2Xla_has_vjp W b) ε hε γ β

Instances For

theorem Proofs.dwbsB_differentiable (N : ℕ) {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ℝ) (hε : 0 < ε) (γ β : Vec c) :

Differentiable ℝ (dwbsB N W b ε γ β)

noncomputable def Proofs.dwbsB_has_vjp (N : ℕ) {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ℝ) (hε : 0 < ε) (γ β : Vec c) :

HasVJP (dwbsB N W b ε γ β)

Equations

Proofs.dwbsB_has_vjp N W b ε hε γ β = Proofs.bnSwishStage_has_vjp N (Proofs.depthwiseFlat W b) ⋯ (Proofs.depthwiseFlat_has_vjp W b) ε hε γ β

Instances For

theorem Proofs.dwbsSB_differentiable (N : ℕ) {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ℝ) (hε : 0 < ε) (γ β : Vec c) :

Differentiable ℝ (dwbsSB N W b ε γ β)

noncomputable def Proofs.dwbsSB_has_vjp (N : ℕ) {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ℝ) (hε : 0 < ε) (γ β : Vec c) :

HasVJP (dwbsSB N W b ε γ β)

Equations

Proofs.dwbsSB_has_vjp N W b ε hε γ β = Proofs.bnSwishStage_has_vjp N (Proofs.depthwiseStride2Flat W b) ⋯ (Proofs.depthwiseStride2Flat_has_vjp W b) ε hε γ β

Instances For

theorem Proofs.seB_differentiable (N : ℕ) {c h w r : ℕ} (W₁ : Mat c r) (b₁ : Vec r) (W₂ : Mat r c) (b₂ : Vec c) :

Differentiable ℝ (seB N W₁ b₁ W₂ b₂)

noncomputable def Proofs.seB_has_vjp (N : ℕ) {c h w r : ℕ} (W₁ : Mat c r) (b₁ : Vec r) (W₂ : Mat r c) (b₂ : Vec c) :

HasVJP (seB N W₁ b₁ W₂ b₂)

Equations

Proofs.seB_has_vjp N W₁ b₁ W₂ b₂ = Proofs.batchMap_has_vjp (Proofs.seBlockFull W₁ b₁ W₂ b₂) (Proofs.seBlockFull_has_vjp W₁ b₁ W₂ b₂) ⋯

Instances For

theorem Proofs.projB_differentiable (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

Differentiable ℝ (projB N W b ε γ β)

noncomputable def Proofs.projB_has_vjp (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

HasVJP (projB N W b ε γ β)

Equations

Proofs.projB_has_vjp N W b ε hε γ β = Proofs.bnStage_has_vjp N (Proofs.flatConv W b) ⋯ (Proofs.flatConv_has_vjp W b) ε hε γ β

Instances For