Documentation

LeanMlir.Proofs.Foundation.Bf16GradNodes

The bf16 gradient nodes — every *GradBBf16 kind the suite emits, folded once #

A bf16 render does NOT consume the f32 gradient node. It emits its own *GradBBf16 constructor, whose den rounds the operands going in and — for every kind but one — rounds the result ONCE, outside the batch sum: the emitted convolution contracts the batch inside a single op and stores its bf16 result once, so a rounding per summand would claim a coarser computation than the hardware performs. That is a different real number from the f32 node's, and it needs its own certificate: the bf16 twins do not consume the f32 node.

This file is the whole bf16 op table, stated per op kind: each lemma says the node denotes the certified Σ_n gradient at the ROUNDED operands, rounded. The proofs are the f32 fold's — congrArg rnd peels the outer rounding, Finset.sum_congr the batch, and the inner equality is the per-example certificate at rounded slices.

kindlemmaf32 peeremitted by
convWeightGradBBf16convWGradBBf16_denGradNodeB.convWGradB_denevery net
convStridedWeightGradBBf16 (symmetric)convStridedWGradBBf16_denGradNodeB.convStridedWGradB_denResNet-34/50, MobileNetV4's fused stage, ConvNeXt's downsamples
convStridedXlaWeightGradBBf16 (XLA-SAME)convStridedXlaWGradBBf16_denGradNodeB.convStridedXlaWGradB_denEfficientNet-B0's, MobileNetV2's and MobileNetV4's stems
convStride4WeightGradBBf16convStride4WGradBBf16_denGradNodeB.psWGradB_denConvNeXt's patchify stem
depthwiseWeightGradBBf16depthwiseWGradBBf16_denGradNodeB.depthwiseWGradB_denB0, MobileNetV2, MobileNetV4, ConvNeXt
depthwiseStridedWeightGradBBf16 (symmetric)depthwiseStridedWGradBBf16_denGradNodeB.depthwiseStridedWGradB_denB0, MobileNetV4
depthwiseStridedXlaWeightGradBBf16 (XLA-SAME)depthwiseStridedXlaWGradBBf16_denGradNodeB.depthwiseStridedXlaWGradB_denMobileNetV2
rowDenseWeightGradBBf16rowDenseWGradBBf16_denViTFoldGB.rowDenseWeightGradB_denViT's Q/K/V/O and MLP denses
patchEmbedWeightGradBBf16patchEmbedWGradBBf16_denViTFoldGB.patchEmbedWeightGradB_denViT's patch embed

Note: rowDenseWeightGradBBf16 has NO outer rounding, and that is the measurement rather than an omission: its dot_general contracts batch and token in one op and keeps its f32-typed result deliberately (StableHLO.Basic's constructor says why), so only the two leaf reads round.

Note: Padding rides along invisibly, as in the f32 folds: the symmetric and XLA-SAME strided kinds have identical types and identical emitted shapes, and only the certificate tells them apart.

Note: BatchNorm, LayerNorm and the dense head have no bf16 twin here or anywhere: every bf16 net in the suite keeps them in f32, so their γ/β and weight nodes are the f32 folds' in both worlds.

theorem Proofs.Bf16Fold.convWGradBBf16_den {N ic oc h w kH kW : ℕ} (rnd : ℝ → ℝ) (xN cotN : String) (b : Vec oc) (x : Vec (N * (ic * h * w))) (W : Kernel4 oc ic kH kW) (cot : Vec (N * (oc * h * w))) (idx : Fin (oc * ic * kH * kW)) :
StableHLO.den (StableHLO.SHlo.convWeightGradBBf16 rnd xN b x W (StableHLO.SHlo.operand cotN cot)) idx = rnd (∑ n : Fin N, ∑ j : Fin (oc * h * w), pdiv (fun (v' : Vec (oc * ic * kH * kW)) => (conv2d (Kernel4.unflatten v') b (Tensor3.unflatten fun (j : Fin (ic * h * w)) => rnd (StableHLO.batchSlice N (ic * h * w) x n j))).flatten) W.flatten idx j * rnd (StableHLO.batchSlice N (oc * h * w) cot n j))

bf16 conv weight GRADIENT denotes the certified Σ_n weight gradient at the rounded operands, rounded once. Every 1×1 and 3×3 in every bf16 artifact.

theorem Proofs.Bf16Fold.convStridedWGradBBf16_den {N ic oc h w kH kW : ℕ} (rnd : ℝ → ℝ) (xN cotN : String) (b : Vec oc) (x : Vec (N * (ic * (2 * h) * (2 * w)))) (W : Kernel4 oc ic kH kW) (cot : Vec (N * (oc * h * w))) (idx : Fin (oc * ic * kH * kW)) :
StableHLO.den (StableHLO.SHlo.convStridedWeightGradBBf16 rnd xN b x W (StableHLO.SHlo.operand cotN cot)) idx = rnd (∑ n : Fin N, ∑ j : Fin (oc * h * w), pdiv (fun (v' : Vec (oc * ic * kH * kW)) => flatConvStride2 (Kernel4.unflatten v') b fun (j : Fin (ic * (2 * h) * (2 * w))) => rnd (StableHLO.batchSlice N (ic * (2 * h) * (2 * w)) x n j)) W.flatten idx j * rnd (StableHLO.batchSlice N (oc * h * w) cot n j))

bf16 SYMMETRIC strided conv weight GRADIENT, rounded once. ResNet's downsamples and 7×7 stem, MobileNetV4's fused 3×3/s2, ConvNeXt's 2×2/s2 downsamples.

theorem Proofs.Bf16Fold.convStridedXlaWGradBBf16_den {N ic oc h w kH kW : ℕ} (rnd : ℝ → ℝ) (xN cotN : String) (b : Vec oc) (x : Vec (N * (ic * (2 * h) * (2 * w)))) (W : Kernel4 oc ic kH kW) (cot : Vec (N * (oc * h * w))) (idx : Fin (oc * ic * kH * kW)) :
StableHLO.den (StableHLO.SHlo.convStridedXlaWeightGradBBf16 rnd xN b x W (StableHLO.SHlo.operand cotN cot)) idx = rnd (∑ n : Fin N, ∑ j : Fin (oc * h * w), pdiv (fun (v' : Vec (oc * ic * kH * kW)) => flatConvStride2Xla (Kernel4.unflatten v') b fun (j : Fin (ic * (2 * h) * (2 * w))) => rnd (StableHLO.batchSlice N (ic * (2 * h) * (2 * w)) x n j)) W.flatten idx j * rnd (StableHLO.batchSlice N (oc * h * w) cot n j))

bf16 XLA-SAME strided conv weight GRADIENT, rounded once. The TF-origin stems (EfficientNet-B0, MobileNetV2, MobileNetV4). Note: flatConvStride2Xla, not flatConvStride2.

theorem Proofs.Bf16Fold.convStride4WGradBBf16_den {N ic oc h w kH kW : ℕ} (rnd : ℝ → ℝ) (xN cotN : String) (b : Vec oc) (x : Vec (N * (ic * (2 * (2 * h)) * (2 * (2 * w))))) (W : Kernel4 oc ic kH kW) (cot : Vec (N * (oc * h * w))) (idx : Fin (oc * ic * kH * kW)) :
StableHLO.den (StableHLO.SHlo.convStride4WeightGradBBf16 rnd xN b x W (StableHLO.SHlo.operand cotN cot)) idx = rnd (∑ n : Fin N, ∑ j : Fin (oc * h * w), pdiv (fun (v' : Vec (oc * ic * kH * kW)) => flatConvStride4 (Kernel4.unflatten v') b fun (j : Fin (ic * (2 * (2 * h)) * (2 * (2 * w)))) => rnd (StableHLO.batchSlice N (ic * (2 * (2 * h)) * (2 * (2 * w))) x n j)) W.flatten idx j * rnd (StableHLO.batchSlice N (oc * h * w) cot n j))

bf16 4×4/s4 patchify-stem weight GRADIENT, rounded once. ConvNeXt's stem.

theorem Proofs.Bf16Fold.depthwiseWGradBBf16_den {N c h w kH kW : ℕ} (rnd : ℝ → ℝ) (xN cotN : String) (b : Vec c) (x : Vec (N * (c * h * w))) (W : DepthwiseKernel c kH kW) (cot : Vec (N * (c * h * w))) (idx : Fin (c * kH * kW)) :
StableHLO.den (StableHLO.SHlo.depthwiseWeightGradBBf16 rnd xN b x W (StableHLO.SHlo.operand cotN cot)) idx = rnd (∑ n : Fin N, ∑ j : Fin (c * h * w), pdiv (fun (v' : Vec (c * kH * kW)) => (depthwiseConv2d (Tensor3.unflatten v') b (Tensor3.unflatten fun (j : Fin (c * h * w)) => rnd (StableHLO.batchSlice N (c * h * w) x n j))).flatten) (Tensor3.flatten W) idx j * rnd (StableHLO.batchSlice N (c * h * w) cot n j))

bf16 depthwise weight GRADIENT, rounded once. Every stride-1 depthwise.

theorem Proofs.Bf16Fold.depthwiseStridedWGradBBf16_den {N c h w kH kW : ℕ} (rnd : ℝ → ℝ) (xN cotN : String) (b : Vec c) (x : Vec (N * (c * (2 * h) * (2 * w)))) (W : DepthwiseKernel c kH kW) (cot : Vec (N * (c * h * w))) (idx : Fin (c * kH * kW)) :
StableHLO.den (StableHLO.SHlo.depthwiseStridedWeightGradBBf16 rnd xN b x W (StableHLO.SHlo.operand cotN cot)) idx = rnd (∑ n : Fin N, ∑ j : Fin (c * h * w), pdiv (fun (v' : Vec (c * kH * kW)) => depthwiseStride2Flat (Tensor3.unflatten v') b fun (j : Fin (c * (2 * h) * (2 * w))) => rnd (StableHLO.batchSlice N (c * (2 * h) * (2 * w)) x n j)) (Tensor3.flatten W) idx j * rnd (StableHLO.batchSlice N (c * h * w) cot n j))

bf16 SYMMETRIC strided depthwise weight GRADIENT, rounded once. B0's stride-2 MBConvs and MobileNetV4's rows 1, 3, 11.

theorem Proofs.Bf16Fold.depthwiseStridedXlaWGradBBf16_den {N c h w kH kW : ℕ} (rnd : ℝ → ℝ) (xN cotN : String) (b : Vec c) (x : Vec (N * (c * (2 * h) * (2 * w)))) (W : DepthwiseKernel c kH kW) (cot : Vec (N * (c * h * w))) (idx : Fin (c * kH * kW)) :
StableHLO.den (StableHLO.SHlo.depthwiseStridedXlaWeightGradBBf16 rnd xN b x W (StableHLO.SHlo.operand cotN cot)) idx = rnd (∑ n : Fin N, ∑ j : Fin (c * h * w), pdiv (fun (v' : Vec (c * kH * kW)) => depthwiseStride2FlatXla (Tensor3.unflatten v') b fun (j : Fin (c * (2 * h) * (2 * w))) => rnd (StableHLO.batchSlice N (c * (2 * h) * (2 * w)) x n j)) (Tensor3.flatten W) idx j * rnd (StableHLO.batchSlice N (c * h * w) cot n j))

bf16 XLA-SAME strided depthwise weight GRADIENT, rounded once. MobileNetV2's four stride-2 depthwises, and no other net's. Note: depthwiseStride2FlatXla.

theorem Proofs.Bf16Fold.rowDenseWGradBBf16_den {N tk a c : ℕ} (rnd : ℝ → ℝ) (xN cotN : String) (bb : Vec c) (x : Vec (N * (tk * a))) (W : Mat a c) (dy : Vec (N * (tk * c))) (i : Fin a) (j : Fin c) :
StableHLO.den (StableHLO.SHlo.rowDenseWeightGradBBf16 rnd xN x (StableHLO.SHlo.operand cotN dy)) (finProdFinEquiv (i, j)) = ∑ n : Fin N, ∑ o : Fin (tk * c), pdiv (fun (v : Vec (a * c)) => Mat.flatten fun (r : Fin tk) => dense (Mat.unflatten v) bb fun (k : Fin a) => rnd (Mat.unflatten (StableHLO.batchSlice N (tk * a) x n) r k)) W.flatten (finProdFinEquiv (i, j)) o * rnd (StableHLO.batchSlice N (tk * c) dy n o)

bf16 per-token dense weight GRADIENT at the rounded operands — with NO outer rounding: the emitted dot_general keeps its f32-typed result, so only the two leaf reads round. ViT's Wq/Wk/Wv/Wo/Wfc1/Wfc2.

theorem Proofs.Bf16Fold.patchEmbedWGradBBf16_den {ic H W P tk D N : ℕ} (rnd : ℝ → ℝ) (xN cotN : String) (bc cls : Vec D) (pos : Mat (tk + 1) D) (img : Vec (N * (ic * H * W))) (Wp : Kernel4 D ic P P) (dy : Vec (N * ((tk + 1) * D))) (d : Fin D) (c : Fin ic) (kh kw : Fin P) :
StableHLO.den (StableHLO.SHlo.patchEmbedWeightGradBBf16 rnd xN img (StableHLO.SHlo.operand cotN dy)) (finProdFinEquiv (finProdFinEquiv (finProdFinEquiv (d, c), kh), kw)) = rnd (∑ n : Fin N, ∑ o : Fin ((tk + 1) * D), pdiv (fun (v : Vec (D * ic * P * P)) => patchEmbedFlat ic H W P tk D (Kernel4.unflatten v) bc cls pos fun (j : Fin (ic * H * W)) => rnd (StableHLO.batchSlice N (ic * H * W) img n j)) Wp.flatten (finProdFinEquiv (finProdFinEquiv (finProdFinEquiv (d, c), kh), kw)) o * rnd (StableHLO.batchSlice N ((tk + 1) * D) dy n o))

bf16 patch-embed conv weight GRADIENT at the rounded operands, rounded once. ViT's 16×16/s16 stem; the outer rnd wraps the whole batch sum because the emit contracts the batch inside one bf16-typed convolution.