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 consume the same node" was written in three fold headers and is false.

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 — congr 1 peels the outer rounding, Finset.sum_congr the batch, and the inner equality is the per-example certificate at rounded slices.

kindf32 peeremitted by
convWeightGradBBf16ResNet34PoCB.convWGradB_denevery net
convStridedWeightGradBBf16 (symmetric)ResNet34PoCB.convStridedWGradB_denResNet-34/50, MobileNetV4's fused stage, ConvNeXt's downsamples
convStridedXlaWeightGradBBf16 (XLA-SAME)EnetPoCG.convStridedXlaWGradB_denEfficientNet-B0's, MobileNetV2's and MobileNetV4's stems
convStride4WeightGradBBf16CnxPoCGB.psWGradB_denConvNeXt's patchify stem
depthwiseWeightGradBBf16EnetPoCG.depthwiseWGradB_denB0, MobileNetV2, MobileNetV4, ConvNeXt
depthwiseStridedWeightGradBBf16 (symmetric)EnetPoCG.depthwiseStridedWGradB_denB0, MobileNetV4
depthwiseStridedXlaWeightGradBBf16 (XLA-SAME)Mnv2PaperPoCG.depthwiseStridedXlaWGradB_denMobileNetV2
rowDenseWeightGradBBf16ViTPoCGB.rowDenseWeightGradB_denViT's Q/K/V/O and MLP denses
patchEmbedWeightGradBBf16ViTPoCGB.patchEmbedWeightGradB_denViT's patch embed

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.lean's constructor says why), so only the two leaf reads round.

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

⛔ 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.Bf16PoC.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.Bf16PoC.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.Bf16PoC.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). ⚠ flatConvStride2Xla, not flatConvStride2.

theorem Proofs.Bf16PoC.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.Bf16PoC.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.Bf16PoC.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.Bf16PoC.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. ⚠ depthwiseStride2FlatXla.

theorem Proofs.Bf16PoC.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.Bf16PoC.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)) => patchEmbed_flat 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.