Per-channel BatchNorm (Chapter 5 Milestone B8) — the block-diagonal VJP

Chapters 4–5 used a per-example global BatchNorm: one scalar (γ, β) over the whole oc·h·w activation (LayerNorm-shaped). Real ResNet wants per-channel BN: normalize each channel-slice independently with its own (γ_c, β_c), γ/β : Vec oc.

Because each channel is independent, the whole Jacobian is block-diagonal across the channel axis — the genuinely-new piece. We get it for free by generalizing the existing rowwise_has_vjp_mat (Tensor.lean, multi-head attention) from a single per-row map to a per-row family g : Fin m → (Vec n → Vec p): viewing the activation as Mat oc (h·w) (row = channel), per-channel BN is exactly fun A => fun c => bnForward (h·w) ε (γ c) (β c) (A c). Its VJP runs each channel's bn_has_vjp on that channel's cotangent slice; the cross-channel blocks vanish.

Everything closes under [propext, Classical.choice, Quot.sound].

theorem Proofs.pdivMat_rowIndep_perRow {m n p : ℕ} (g : Fin m → Vec n → Vec p) (h_g_diff : ∀ (r : Fin m), Differentiable ℝ (g r)) (A : Mat m n) (i : Fin m) (j : Fin n) (k : Fin m) (l : Fin p) : pdivMat (fun (M : Mat m n) (r : Fin m) => g r (M r)) A i j k l = if i = k then pdiv (g k) (A k) j l else 0

Block-diagonal Jacobian of a per-row family. Applying a different map g r to each row r of a matrix keeps the matrix Jacobian block-diagonal across the row axis: output row k depends only on input row k (via g k). The per-row generalization of pdivMat_rowIndep (which fixes one g for all rows).

noncomputable def Proofs.rowwisePerRow_has_vjp_mat {m n p : ℕ} (g : Fin m → Vec n → Vec p) (hg : (r : Fin m) → HasVJP (g r)) (hg_diff : ∀ (r : Fin m), Differentiable ℝ (g r)) : HasVJPMat fun (A : Mat m n) (r : Fin m) => g r (A r)

Row-wise lifting of a per-row HasVJP family. Each row r gets its own map g r (with its own VJP); the matrix backward runs (g r).backward on row r's cotangent. The per-row peer of rowwise_has_vjp_mat.

Equations

• Proofs.rowwisePerRow_has_vjp_mat g hg hg_diff = { backward := fun (A : Proofs.Mat m n) (dY : Proofs.Mat m p) (r : Fin m) (c : Fin n) => (hg r).backward (A r) (dY r) c, correct := ⋯ }

theorem Proofs.rowwisePerRow_flat_differentiable {m n p : ℕ} (g : Fin m → Vec n → Vec p) (h_g_diff : ∀ (r : Fin m), Differentiable ℝ (g r)) : Differentiable ℝ fun (v : Vec (m * n)) => Mat.flatten ((fun (A : Mat m n) (r : Fin m) => g r (A r)) (Mat.unflatten v))

A per-row family flattens to a differentiable Vec → Vec map. The Differentiable witness vjp_comp_at / the network composition needs to thread a per-channel BN through a block.

noncomputable def Proofs.bnPerChannelMat (oc m : ℕ) (ε : ℝ) (γ β : Vec oc) : Mat oc m → Mat oc m

Per-channel BatchNorm (matrix view): BN each row (= channel-slice of m = h·w spatial cells) with its own (γ_c, β_c). The real-ResNet BN that bnForward's global scalar version approximates.

Equations

• Proofs.bnPerChannelMat oc m ε γ β A c = Proofs.bnForward m ε (γ c) (β c) (A c)

noncomputable def Proofs.bnPerChannelMat_has_vjp (oc m : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) : HasVJPMat (bnPerChannelMat oc m ε γ β)

Per-channel BN VJP (block-diagonal). Each channel runs its own bn_has_vjp; the cross-channel Jacobian blocks vanish (pdivMat_rowIndep_perRow).

Equations

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

noncomputable def Proofs.bnPerChannelFlat (oc m : ℕ) (ε : ℝ) (γ β : Vec oc) : Vec (oc * m) → Vec (oc * m)

Per-channel BN as a flat-vector op Vec (oc·m) → Vec (oc·m) (row-major, channel c = the m-wide slab at flat positions finProdFinEquiv (c, ·)).

Equations

• Proofs.bnPerChannelFlat oc m ε γ β v = (Proofs.bnPerChannelMat oc m ε γ β (Proofs.Mat.unflatten v)).flatten

noncomputable def Proofs.bnPerChannelFlat_has_vjp (oc m : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) : HasVJP (bnPerChannelFlat oc m ε γ β)

Per-channel BN flat VJP — the block-diagonal matrix VJP bridged to Vec.

Equations

• Proofs.bnPerChannelFlat_has_vjp oc m ε hε γ β = Proofs.hasVJPMat_to_hasVJP (Proofs.bnPerChannelMat_has_vjp oc m ε hε γ β)

theorem Proofs.bnPerChannelFlat_differentiable (oc m : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) : Differentiable ℝ (bnPerChannelFlat oc m ε γ β)

Per-channel BN is differentiable everywhere (ε > 0). The composition witness.

theorem Proofs.bnPerChannelFlat_has_vjp_correct (oc m : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (v dy : Vec (oc * m)) (i : Fin (oc * m)) : (bnPerChannelFlat_has_vjp oc m ε hε γ β).backward v dy i = ∑ j : Fin (oc * m), pdiv (bnPerChannelFlat oc m ε γ β) v i j * dy j

Per-channel BN VJP correctness (ℝ-headline): the flat backward equals the pdiv-contracted (block-diagonal) Jacobian of per-channel BN.

noncomputable def Proofs.bnPerChannel_grad_input (oc m : ℕ) (ε : ℝ) (γ : Vec oc) (x dy : Vec (oc * m)) : Vec (oc * m)

Per-channel consolidated BN input-gradient — the renderable closed form: run the per-example three-term bn_grad_input on each channel-slice (m = h·w spatial cells), reusing that channel's γ_c. This is exactly what a bnPerChannelBack SHlo op / renderLNBack-per-channel emits; the abstract bnPerChannelFlat_has_vjp.backward is the spec it must match.

Equations

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

theorem Proofs.bnPerChannel_grad_input_correct (oc m : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x dy : Vec (oc * m)) (i : Fin (oc * m)) : bnPerChannel_grad_input oc m ε γ x dy i = ∑ j : Fin (oc * m), pdiv (bnPerChannelFlat oc m ε γ β) x i j * dy j

Renderable backward is faithful (ℝ-headline): the per-channel consolidated gradient equals the pdiv-contracted Jacobian of per-channel BN, under 0 < ε. Each channel reduces to the per-example bn_input_grad_correct. The licence to render per-channel BN's backward as the three-term formula per channel.

The network carries its activations in the Tensor3 flat layout (oc*h)*w (flatConv etc.), but bnPerChannelFlat is defined on the Mat-split layout oc*(h*w) (row c = the h·w spatial cells of channel c). The two Vecs have the same size; they differ only in how finProdFinEquiv associates the product. So the bridge is a pure re-association reindex (a permutation of coordinates) — a reindexCLM whose VJP is the scatter pdiv_reindex gives, exactly like decimateFlat. Conjugating bnPerChannelFlat by this bridge yields per-channel BN acting on the network's Tensor3 activations, with its VJP for free via vjp_comp.

noncomputable def Proofs.reassocFwdIdx (oc h w : ℕ) (mIdx : Fin (oc * (h * w))) : Fin (oc * h * w)

Re-association index Fin (oc*(h*w)) → Fin (oc*h*w): a Mat-split flat index (c, s) with s ↔ (hi, wi) maps to the Tensor3 flat index ((c, hi), wi). A pure product re-association — no arithmetic.

Equations

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

noncomputable def Proofs.reassocBackIdx (oc h w : ℕ) (t : Fin (oc * h * w)) : Fin (oc * (h * w))

Re-association index Fin (oc*h*w) → Fin (oc*(h*w)): the inverse direction, Tensor3 ((c, hi), wi) ↦ Mat-split (c, (hi, wi)).

Equations

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

theorem Proofs.reassocFwdIdx_reassocBackIdx (oc h w : ℕ) (t : Fin (oc * h * w)) : reassocFwdIdx oc h w (reassocBackIdx oc h w t) = t

The two re-association indices are mutual inverses — the bridge is a genuine relabeling (so conjugating by it really is per-channel BN, just in Tensor3 coordinates). Pure finProdFinEquiv round-trip.

theorem Proofs.reassocBackIdx_reassocFwdIdx (oc h w : ℕ) (mIdx : Fin (oc * (h * w))) : reassocBackIdx oc h w (reassocFwdIdx oc h w mIdx) = mIdx

noncomputable def Proofs.reassocFwd (oc h w : ℕ) : Vec (oc * h * w) → Vec (oc * (h * w))

Tensor3 → Mat-split reindex: read the ((c,hi),wi) cell at Mat position (c, (hi,wi)). A reindexCLM, hence continuous-linear / differentiable.

Equations

• Proofs.reassocFwd oc h w y k = y (Proofs.reassocFwdIdx oc h w k)

noncomputable def Proofs.reassocBack (oc h w : ℕ) : Vec (oc * (h * w)) → Vec (oc * h * w)

Mat-split → Tensor3 reindex (the inverse relabeling).

Equations

• Proofs.reassocBack oc h w y k = y (Proofs.reassocBackIdx oc h w k)

Training BN computes μ/σ² from the activation it is normalizing; inference BN consumes frozen statistics (the driver's EMA'd running mean/var) and is therefore a plain affine map — pointwise in the activation, with no reduction at all. That is exactly why an eval forward built on it is class-batch-independent: an example's logits do not depend on which other examples share its batch.

The chain below mirrors bnForward → bnPerChannelMat → bnPerChannelFlat → bnPerChannelTensor3 one-for-one, so the eval op drops into the same layout bridge as the training op.

noncomputable def Proofs.bnEvalForward (m : ℕ) (ε γ β μ v : ℝ) (x : Vec m) : Vec m

Inference BN on one channel's m activations: yᵢ = γ · (xᵢ − μ) · (var + ε)^(−1/2) + β, with μ/var supplied rather than computed from x. The bnForward peer — note it takes x pointwise, where bnForward reduces over all of x to get its own μ/σ².

Equations

• Proofs.bnEvalForward m ε γ β μ v x i = γ * ((x i - μ) * (1 / √(v + ε))) + β

noncomputable def Proofs.bnPerChannelEvalMat (oc m : ℕ) (ε : ℝ) (γ β μ v : Vec oc) : Mat oc m → Mat oc m

Per-channel inference BN (Mat layout) — row c gets channel c's frozen stats.

Equations

• Proofs.bnPerChannelEvalMat oc m ε γ β μ v A c = Proofs.bnEvalForward m ε (γ c) (β c) (μ c) (v c) (A c)

noncomputable def Proofs.bnPerChannelEvalFlat (oc m : ℕ) (ε : ℝ) (γ β μ v : Vec oc) : Vec (oc * m) → Vec (oc * m)

Per-channel inference BN (flat layout) — the bnPerChannelFlat peer.

Equations

• Proofs.bnPerChannelEvalFlat oc m ε γ β μ v x = (Proofs.bnPerChannelEvalMat oc m ε γ β μ v (Proofs.Mat.unflatten x)).flatten

noncomputable def Proofs.bnPerChannelEvalTensor3 (oc h w : ℕ) (ε : ℝ) (γ β μ v : Vec oc) : Vec (oc * h * w) → Vec (oc * h * w)

Per-channel inference BN (Tensor3 layout) — the bnPerChannelTensor3 peer, through the same reassoc bridge. This is what SHlo.bnPerChannelEvalF denotes.

Equations

• Proofs.bnPerChannelEvalTensor3 oc h w ε γ β μ v = Proofs.reassocBack oc h w ∘ Proofs.bnPerChannelEvalFlat oc (h * w) ε γ β μ v ∘ Proofs.reassocFwd oc h w

theorem Proofs.bnEvalForward_differentiable (m : ℕ) (ε γ β μ v : ℝ) : Differentiable ℝ (bnEvalForward m ε γ β μ v)

Inference BN is differentiable everywhere — it is affine in x, so unlike the training BN this needs no 0 < ε hypothesis: ε only enters the constant scale factor.

noncomputable def Proofs.bnPerChannel_grad_gamma (oc m : ℕ) (ε : ℝ) (v dy : Vec (oc * m)) : Vec oc

The rendered per-channel γ gradient: dγ_c = Σ_{s} dy_(c,s) · x̂_(c,s) (the reduce over batch/spatial of dy·x̂ in cifarBnTrainStepStructured's bnParamGradPC). x̂ is recomputed from the saved BN input v (the conv output). Lives here (not CifarBnClose) so the bnGammaSgd SHlo op's den can reference it.

Equations

• Proofs.bnPerChannel_grad_gamma oc m ε v dy c = ∑ s : Fin m, dy (finProdFinEquiv (c, s)) * Proofs.bnXhat m ε (Proofs.Mat.unflatten v c) s

noncomputable def Proofs.bnPerChannel_grad_beta (oc m : ℕ) (dy : Vec (oc * m)) : Vec oc

The rendered per-channel β gradient: dβ_c = Σ_{s} dy_(c,s).

Equations

• Proofs.bnPerChannel_grad_beta oc m dy c = ∑ s : Fin m, dy (finProdFinEquiv (c, s))

theorem Proofs.reassocFwd_differentiable (oc h w : ℕ) : Differentiable ℝ (reassocFwd oc h w)

theorem Proofs.reassocBack_differentiable (oc h w : ℕ) : Differentiable ℝ (reassocBack oc h w)

noncomputable def Proofs.reassocFwd_has_vjp (oc h w : ℕ) : HasVJP (reassocFwd oc h w)

VJP of the forward reindex — the scatter pdiv_reindex gives. Mirrors decimateFlat_has_vjp.

Equations

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

noncomputable def Proofs.reassocBack_has_vjp (oc h w : ℕ) : HasVJP (reassocBack oc h w)

Equations

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

theorem Proofs.reassocBack_has_vjp_backward_eq (oc h w : ℕ) (v : Vec (oc * (h * w))) (dy : Vec (oc * h * w)) : (reassocBack_has_vjp oc h w).backward v dy = reassocFwd oc h w dy

The bridge is a permutation, so each reindex's VJP backward is just the inverse reindex (the single matching delta survives the scatter). These two collapse the vjp_comp backwards into a clean closed form for bnPerChannelTensor3.

theorem Proofs.reassocFwd_has_vjp_backward_eq (oc h w : ℕ) (v : Vec (oc * h * w)) (dy : Vec (oc * (h * w))) : (reassocFwd_has_vjp oc h w).backward v dy = reassocBack oc h w dy

noncomputable def Proofs.bnPerChannelTensor3 (oc h w : ℕ) (ε : ℝ) (γ β : Vec oc) : Vec (oc * h * w) → Vec (oc * h * w)

Per-channel BatchNorm on the Tensor3 (oc*h)*w activation layout. Conjugate the Mat-split bnPerChannelFlat by the layout bridge: relabel to Mat-split, normalize each channel over its h·w spatial cells, relabel back. This is the op B9 wires into the ResNet-34 trainer (its den target).

Equations

• Proofs.bnPerChannelTensor3 oc h w ε γ β = Proofs.reassocBack oc h w ∘ Proofs.bnPerChannelFlat oc (h * w) ε γ β ∘ Proofs.reassocFwd oc h w

theorem Proofs.bnPerChannelTensor3_differentiable (oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) : Differentiable ℝ (bnPerChannelTensor3 oc h w ε γ β)

noncomputable def Proofs.bnPerChannelTensor3_has_vjp (oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) : HasVJP (bnPerChannelTensor3 oc h w ε γ β)

Per-channel BN (Tensor3 layout) VJP — block-diagonal across channels, lifted through the layout bridge by vjp_comp (twice).

Equations

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

theorem Proofs.bnPerChannelTensor3_has_vjp_correct (oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x dy : Vec (oc * h * w)) (i : Fin (oc * h * w)) : (bnPerChannelTensor3_has_vjp oc h w ε hε γ β).backward x dy i = ∑ j : Fin (oc * h * w), pdiv (bnPerChannelTensor3 oc h w ε γ β) x i j * dy j

Per-channel BN (Tensor3 layout) VJP correctness (ℝ-headline): the backward equals the pdiv-contracted (block-diagonal) Jacobian of per-channel BN on the network's activation layout. The licence to wire per-channel BN into ResNet-34.

theorem Proofs.bnPerChannelTensor3_has_vjp_backward_eq (oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x dy : Vec (oc * h * w)) : (bnPerChannelTensor3_has_vjp oc h w ε hε γ β).backward x dy = reassocBack oc h w ((bnPerChannelFlat_has_vjp oc (h * w) ε hε γ β).backward (reassocFwd oc h w x) (reassocFwd oc h w dy))

The composed vjp_comp backward collapses (the bridge reindexes are permutations): per-channel BN's Tensor3 backward is the Mat-split block-diagonal backward, conjugated by the layout bridge.

noncomputable def Proofs.bnPerChannelTensor3_grad_input (oc h w : ℕ) (ε : ℝ) (γ : Vec oc) (x dy : Vec (oc * h * w)) : Vec (oc * h * w)

Renderable per-channel BN backward on the Tensor3 (oc*h)*w layout — relabel to Mat-split, run the per-channel consolidated three-term bnPerChannel_grad_input, relabel back. This is exactly what the bnPerChannelBack SHlo op emits (per-channel renderLNBack, reducing over the spatial axis); its faithfulness spec is below.

Equations

• Proofs.bnPerChannelTensor3_grad_input oc h w ε γ x dy = Proofs.reassocBack oc h w (Proofs.bnPerChannel_grad_input oc (h * w) ε γ (Proofs.reassocFwd oc h w x) (Proofs.reassocFwd oc h w dy))

theorem Proofs.bnPerChannelTensor3_grad_input_correct (oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x dy : Vec (oc * h * w)) (i : Fin (oc * h * w)) : bnPerChannelTensor3_grad_input oc h w ε γ x dy i = ∑ j : Fin (oc * h * w), pdiv (bnPerChannelTensor3 oc h w ε γ β) x i j * dy j

Renderable Tensor3 backward is faithful (ℝ-headline): equals the pdiv-contracted (block-diagonal) Jacobian of per-channel BN on the network's activation layout, under 0 < ε. The licence to render per-channel BN's backward in ResNet-34.

noncomputable def Proofs.bnchwFwdIdx (N oc h w : ℕ) (mIdx : Fin (oc * (N * (h * w)))) : Fin (N * (oc * (h * w)))

Transpose-reindex Fin (oc*(N*(h*w))) → Fin (N*(oc*(h*w))): a per-channel Mat index (c, (n, s)) maps to the network [N,C,H,W] flat index (n, (c, s)) — swap the batch and channel axes (s ↔ (hi,wi) carried along). A permutation.

Equations

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

noncomputable def Proofs.bnchwBackIdx (N oc h w : ℕ) (t : Fin (N * (oc * (h * w)))) : Fin (oc * (N * (h * w)))

Transpose-reindex Fin (N*(oc*(h*w))) → Fin (oc*(N*(h*w))): the inverse, network (n, (c, s)) ↦ per-channel Mat (c, (n, s)).

Equations

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

theorem Proofs.bnchwFwdIdx_bnchwBackIdx (N oc h w : ℕ) (t : Fin (N * (oc * (h * w)))) : bnchwFwdIdx N oc h w (bnchwBackIdx N oc h w t) = t

theorem Proofs.bnchwBackIdx_bnchwFwdIdx (N oc h w : ℕ) (mIdx : Fin (oc * (N * (h * w)))) : bnchwBackIdx N oc h w (bnchwFwdIdx N oc h w mIdx) = mIdx

noncomputable def Proofs.bnchwFwd (N oc h w : ℕ) : Vec (N * (oc * (h * w))) → Vec (oc * (N * (h * w)))

[N,C,H,W] → [C,N·H·W] reindex (gather the network cell at the Mat position).

Equations

• Proofs.bnchwFwd N oc h w y k = y (Proofs.bnchwFwdIdx N oc h w k)

noncomputable def Proofs.bnchwBack (N oc h w : ℕ) : Vec (oc * (N * (h * w))) → Vec (N * (oc * (h * w)))

[C,N·H·W] → [N,C,H,W] reindex (the inverse relabeling).

Equations

• Proofs.bnchwBack N oc h w y k = y (Proofs.bnchwBackIdx N oc h w k)

theorem Proofs.bnchwFwd_differentiable (N oc h w : ℕ) : Differentiable ℝ (bnchwFwd N oc h w)

theorem Proofs.bnchwBack_differentiable (N oc h w : ℕ) : Differentiable ℝ (bnchwBack N oc h w)

noncomputable def Proofs.bnchwFwd_has_vjp (N oc h w : ℕ) : HasVJP (bnchwFwd N oc h w)

Equations

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

noncomputable def Proofs.bnchwBack_has_vjp (N oc h w : ℕ) : HasVJP (bnchwBack N oc h w)

Equations

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

theorem Proofs.bnchwBack_has_vjp_backward_eq (N oc h w : ℕ) (v : Vec (oc * (N * (h * w)))) (dy : Vec (N * (oc * (h * w)))) : (bnchwBack_has_vjp N oc h w).backward v dy = bnchwFwd N oc h w dy

theorem Proofs.bnchwFwd_has_vjp_backward_eq (N oc h w : ℕ) (v : Vec (N * (oc * (h * w)))) (dy : Vec (oc * (N * (h * w)))) : (bnchwFwd_has_vjp N oc h w).backward v dy = bnchwBack N oc h w dy

noncomputable def Proofs.bnBatchTensor4 (N oc h w : ℕ) (ε : ℝ) (γ β : Vec oc) : Vec (N * (oc * (h * w))) → Vec (N * (oc * (h * w)))

Batch-norm per channel on the network's [N,C,H,W] layout. Conjugate the Mat-split bnPerChannelFlat (with m = N·h·w, the whole batch's cells per channel) by the transpose bridge: relabel [N,C,H,W] → [C, N·H·W], normalize each channel over ALL its batch+spatial cells, relabel back. The EfficientNet normalization.

Equations

• Proofs.bnBatchTensor4 N oc h w ε γ β = Proofs.bnchwBack N oc h w ∘ Proofs.bnPerChannelFlat oc (N * (h * w)) ε γ β ∘ Proofs.bnchwFwd N oc h w

theorem Proofs.bnBatchTensor4_differentiable (N oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) : Differentiable ℝ (bnBatchTensor4 N oc h w ε γ β)

noncomputable def Proofs.bnBatchTensor4_has_vjp (N oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) : HasVJP (bnBatchTensor4 N oc h w ε γ β)

Batch-norm (network layout) VJP — block-diagonal across channels (now coupling the whole batch within each channel), lifted through the transpose bridge.

Equations

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

theorem Proofs.bnBatchTensor4_has_vjp_correct (N oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x dy : Vec (N * (oc * (h * w)))) (i : Fin (N * (oc * (h * w)))) : (bnBatchTensor4_has_vjp N oc h w ε hε γ β).backward x dy i = ∑ j : Fin (N * (oc * (h * w))), pdiv (bnBatchTensor4 N oc h w ε γ β) x i j * dy j

theorem Proofs.bnBatchTensor4_has_vjp_backward_eq (N oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x dy : Vec (N * (oc * (h * w)))) : (bnBatchTensor4_has_vjp N oc h w ε hε γ β).backward x dy = bnchwBack N oc h w ((bnPerChannelFlat_has_vjp oc (N * (h * w)) ε hε γ β).backward (bnchwFwd N oc h w x) (bnchwFwd N oc h w dy))

noncomputable def Proofs.bnBatchTensor4_grad_input (N oc h w : ℕ) (ε : ℝ) (γ : Vec oc) (x dy : Vec (N * (oc * (h * w)))) : Vec (N * (oc * (h * w)))

Renderable batch-norm backward on the [N,C,H,W] layout — relabel to the per-channel Mat, run the consolidated three-term bnPerChannel_grad_input over the whole batch (m = N·h·w), relabel back. Exactly what the batched bnBatch StableHLO fragment emits (reduce over [0,2,3] per channel).

Equations

• Proofs.bnBatchTensor4_grad_input N oc h w ε γ x dy = Proofs.bnchwBack N oc h w (Proofs.bnPerChannel_grad_input oc (N * (h * w)) ε γ (Proofs.bnchwFwd N oc h w x) (Proofs.bnchwFwd N oc h w dy))

theorem Proofs.bnBatchTensor4_grad_input_correct (N oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) (x dy : Vec (N * (oc * (h * w)))) (i : Fin (N * (oc * (h * w)))) : bnBatchTensor4_grad_input N oc h w ε γ x dy i = ∑ j : Fin (N * (oc * (h * w))), pdiv (bnBatchTensor4 N oc h w ε γ β) x i j * dy j

Renderable batch-norm backward is faithful (ℝ-headline): equals the pdiv-contracted (block-diagonal-across-channels, batch-coupled) Jacobian of batch-norm on the network's [N,C,H,W] layout, under 0 < ε. The licence to render EfficientNet's batch-norm backward as the per-channel three-term formula over the batch.