Documentation

LeanMlir.Proofs.Nets.ResNet.ResNet34SyncB

ResNet-34's data-parallel forward at SYNCHRONISED BatchNorm — replica `r` IS shard `r` #

ResNet34FullB.lean (T2) says the typed batch-BN graph denotes resnet34ForwardB_full N w on one device. ResNet34RenderB renders the data-parallel step differently since 2026-09-21: at replicas > 1 every BatchNorm site is bnFwdSite's sync-BN composition — this replica's mean all-reduced, then Chan's σ²_r + (μ_r − μ)² all-reduced, packed, then bnSyncF — so the statistics each replica normalises by are the GLOBAL batch's. This file is T2's data-parallel twin: that forward graph, stated as a family over the R replicas, denotes on replica r exactly batchShard r of the single-device forward at the global batch R·N.

den (resnet34FwdGraphSync_full R hR N epsStr w e r)
  = batchShard R N nCls (resnet34ForwardB_full (R * N) w X) r

given that each replica's input is its shard of one global batch X. ⭐ The spec does not move: the right-hand side is the committed resnet34ForwardB_full, at N := R·N.

How it is proved #

By induction on the chain, one block at a time, exactly as T2 is — with the shard hypothesis ∀ r, den (e r) = batchShard R N _ X r as the invariant carried from block to block:

every conv, relu, pool, GAP and dense node is a per-example lift, and sharding commutes with it (batchShard_batchMap, and relu pointwise) — den_batchOp_shard, den_relu_shard;
every BatchNorm site is bnSyncSiteLA, whose shard lemma den_bnSyncSiteLA is DataParallelSync.den_bnSyncF_allReduce (P1 on the graph) read at the network's index.

The index seam #

The conv/relu chain runs at the left-assoc index N·(c·h·w); bnSyncF and its statistics nodes at N·(c·(h·w)), the index bnBatchTensor4 is stated at. bnBatchF hides that seam inside its den (bnBatchLA is bnBatchTensor4 conjugated by the mul_assoc relabelling). The sync site cannot, because its statistics subgraph is shared across replicas, so it carries the relabelling as castIdx on the AST value — h ▸ e, the move ConvNeXtRenderB's reassocB already makes. It changes no emitted text: skel never sees an index, and the render writes the same SSA name at both types (bnFwdSite's zbn operand beside its zin one).

What is NOT claimed here #

⚠ The backward and the parameter collectives are the T3 half (ResNet34SyncStepTieB.lean). ⚠ That the R replicas' inputs ARE the shards of one batch is the driver's, as in DataParallelSync.lean. ⚠ The lowerer's all_reduce is trusted as every other op's lowering is.

def Proofs.StableHLO.castIdx {n m : ℕ} (h : n = m) (e : SHlo n) :

Relabel an AST value's index along a proved equality. h ▸ e: the same graph, typed at m instead of n. The emitted text does not change, because skel erases indices.

Equations

Proofs.StableHLO.castIdx h e = h ▸ e

Instances For

theorem Proofs.StableHLO.den_castIdx {n m : ℕ} (h : n = m) (e : SHlo n) :

den (castIdx h e) = fun (i : Fin m) => den e (Fin.cast ⋯ i)

theorem Proofs.StableHLO.laAssoc (N oc h w : ℕ) :

N * (oc * h * w) = N * (oc * (h * w))

The mul_assoc relabelling under N * · — the seam between the network's left-assoc N·(c·h·w) and the BatchNorm ops' N·(c·(h·w)).

theorem Proofs.StableHLO.batchShard_castIdx {R N a b : ℕ} (hab : a = b) (X : Vec (R * N * a)) (r : Fin R) :

batchShard R N b (fun (i : Fin (R * N * b)) => X (Fin.cast ⋯ i)) r = fun (i : Fin (N * b)) => batchShard R N a X r (Fin.cast ⋯ i)

Sharding commutes with relabelling the per-example index. The batch axis is outside the per-example one, so relabelling within an example and cutting the batch do not interact.

theorem Proofs.StableHLO.den_batchOp_shard {R N a b : ℕ} (op : BatchableOp a b) (e : Fin R → SHlo (N * a)) (X : Vec (R * N * a)) (he : ∀ (r : Fin R), den (e r) = batchShard R N a X r) (r : Fin R) :

den (SHlo.batchOp op (e r)) = batchShard R N b (batchMap (R * N) (denOp op) X) r

A per-example node on every replica denotes the shard of the same node on the global batch.

theorem Proofs.StableHLO.den_relu_shard {R N n : ℕ} (e : Fin R → SHlo (N * n)) (X : Vec (R * N * n)) (he : ∀ (r : Fin R), den (e r) = batchShard R N n X r) (r : Fin R) :

den (SHlo.batchOp BatchableOp.relu (e r)) = batchShard R N n (relu (R * N * n) X) r

…stated for relu at the whole-batch relu, the form the committed forward is written in.

theorem Proofs.StableHLO.den_addVB_shard {R N n : ℕ} (a b : Fin R → SHlo (N * n)) (A B : Vec (R * N * n)) (ha : ∀ (r : Fin R), den (a r) = batchShard R N n A r) (hb : ∀ (r : Fin R), den (b r) = batchShard R N n B r) (r : Fin R) :

den ((a r).addVB (b r)) = batchShard R N n (fun (j : Fin (R * N * n)) => A j + B j) r

The residual fan-in on every replica is the shard of the global one.

def Proofs.StableHLO.bnSyncSiteLA (gN bN es t t' : String) (ds ds' : List ℕ) (R : ℕ) (hR : 0 < R) {N oc h w : ℕ} (ε : ℝ) (γ β : Vec oc) (x : Fin R → SHlo (N * (oc * h * w))) (r : Fin R) :

SHlo (N * (oc * h * w))

One sync-BN forward site, at the network index, on replica r — bnFwdSite's replicas > 1 branch: bnSyncF of this replica's operand, reading syncStats over all R replicas' operands (the mean collective t, then Chan's variance collective t'), with the mul_assoc relabelling on the way in and out.

Equations

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

Instances For

theorem Proofs.StableHLO.den_bnSyncSiteLA (gN bN es t t' : String) (ds ds' : List ℕ) (R : ℕ) (hR : 0 < R) {N oc h w : ℕ} (hm : N * (h * w) ≠ 0) (hM : R * N * (h * w) ≠ 0) (ε : ℝ) (γ β : Vec oc) (x : Fin R → SHlo (N * (oc * h * w))) (X : Vec (R * N * (oc * h * w))) (hx : ∀ (r : Fin R), den (x r) = batchShard R N (oc * h * w) X r) (r : Fin R) :

den (bnSyncSiteLA gN bN es t t' ds ds' R hR ε γ β x r) = batchShard R N (oc * h * w) (bnBatchLA (R * N) oc h w ε γ β X) r

⭐⭐ The sync-BN site on replica r is shard r of the global-batch BatchNorm. den_bnSyncF_allReduce (P1 on the graph), carried across the mul_assoc seam: the right-hand side is bnBatchLA — what bnBatchF denotes — at N := R·N.

theorem Proofs.StableHLO.nhw_ne_zero {N h w : ℕ} (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) :

N * (h * w) ≠ 0

The reduction width a BatchNorm site needs nonzero, from the three positive dimensions.

def Proofs.StableHLO.r34IdGraphSync (p epsStr : String) (R : ℕ) (hR : 0 < R) (N h w : ℕ) {c : ℕ} (pw : R34IdW c) (e : Fin R → SHlo (N * (c * h * w))) :

Fin R → SHlo (N * (c * h * w))

Identity basic block at sync-BN, over the replica family: relu(addV(bn₂(conv₂(relu(bn₁(conv₁ e)))), e)) with both BatchNorms bnSyncSiteLA.

Equations

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

Instances For

theorem Proofs.StableHLO.r34IdGraphSync_shard (p epsStr : String) (R : ℕ) (hR : 0 < R) (N h w : ℕ) {c : ℕ} (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (pw : R34IdW c) (e : Fin R → SHlo (N * (c * h * w))) (X : Vec (R * N * (c * h * w))) (he : ∀ (r : Fin R), den (e r) = batchShard R N (c * h * w) X r) (r : Fin R) :

den (r34IdGraphSync p epsStr R hR N h w pw e r) = batchShard R N (c * h * w) (r34IdB (R * N) h w pw X) r

def Proofs.StableHLO.r34DownGraphSync (p epsStr : String) (R : ℕ) (hR : 0 < R) (N h w : ℕ) {ic oc : ℕ} (pw : R34DownW ic oc) (e : Fin R → SHlo (N * (ic * (2 * h) * (2 * w)))) :

Fin R → SHlo (N * (oc * h * w))

Downsample basic block at sync-BN, over the replica family: relu(addV(bnₚ(projection), bn₂(conv₂(relu(bn₁(convStrided₁ e)))))) — projection first, as residualProj and the render order it. Three sync sites.

Equations

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

Instances For

theorem Proofs.StableHLO.r34DownGraphSync_shard (p epsStr : String) (R : ℕ) (hR : 0 < R) (N h w : ℕ) {ic oc : ℕ} (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (pw : R34DownW ic oc) (e : Fin R → SHlo (N * (ic * (2 * h) * (2 * w)))) (X : Vec (R * N * (ic * (2 * h) * (2 * w)))) (he : ∀ (r : Fin R), den (e r) = batchShard R N (ic * (2 * h) * (2 * w)) X r) (r : Fin R) :

den (r34DownGraphSync p epsStr R hR N h w pw e r) = batchShard R N (oc * h * w) (r34DownB (R * N) h w pw X) r

def Proofs.StableHLO.r34StemGraphSync (epsStr : String) (R : ℕ) (hR : 0 < R) (N h w : ℕ) {ic oc : ℕ} (Ws : Kernel4 oc ic 7 7) (bs : Vec oc) (εs : ℝ) (γs βs : Vec oc) (e : Fin R → SHlo (N * (ic * (2 * (2 * h)) * (2 * (2 * w))))) :

Fin R → SHlo (N * (oc * h * w))

Stem at sync-BN, over the replica family: 7×7/s2 conv → sync-BN → relu → 3×3/s2 max-pool.

Equations

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

Instances For

theorem Proofs.StableHLO.r34StemGraphSync_shard (epsStr : String) (R : ℕ) (hR : 0 < R) (N h w : ℕ) {ic oc : ℕ} (hN : 0 < N) (hh : 0 < h) (hw : 0 < w) (Ws : Kernel4 oc ic 7 7) (bs : Vec oc) (εs : ℝ) (γs βs : Vec oc) (e : Fin R → SHlo (N * (ic * (2 * (2 * h)) * (2 * (2 * w))))) (X : Vec (R * N * (ic * (2 * (2 * h)) * (2 * (2 * w))))) (he : ∀ (r : Fin R), den (e r) = batchShard R N (ic * (2 * (2 * h)) * (2 * (2 * w))) X r) (r : Fin R) :

den (r34StemGraphSync epsStr R hR N h w Ws bs εs γs βs e r) = batchShard R N (oc * h * w) (r34StemB (R * N) h w Ws bs εs γs βs X) r

def Proofs.StableHLO.r34HeadGraphSync {R : ℕ} (N h w : ℕ) {c nCls : ℕ} (Wd : Mat c nCls) (bd : Vec nCls) (e : Fin R → SHlo (N * (c * h * w))) :

Fin R → SHlo (N * nCls)

Head over the replica family: GAP then dense — no BatchNorm, so T2's head graph per replica.

Equations

Proofs.StableHLO.r34HeadGraphSync N h w Wd bd e r = Proofs.StableHLO.r34HeadGraphB N h w Wd bd (e r)

Instances For

theorem Proofs.StableHLO.r34HeadGraphSync_shard {R : ℕ} (N h w : ℕ) {c nCls : ℕ} (Wd : Mat c nCls) (bd : Vec nCls) (e : Fin R → SHlo (N * (c * h * w))) (X : Vec (R * N * (c * h * w))) (he : ∀ (r : Fin R), den (e r) = batchShard R N (c * h * w) X r) (r : Fin R) :

den (r34HeadGraphSync N h w Wd bd e r) = batchShard R N nCls (r34HeadB (R * N) h w Wd bd X) r

def Proofs.StableHLO.resnet34FwdGraphSync_full (R : ℕ) (hR : 0 < R) (N : ℕ) (epsStr : String) {nCls : ℕ} (w : R34BWeights nCls) (e : Fin R → SHlo (N * (3 * (2 * (2 * 56)) * (2 * (2 * 56))))) :

Fin R → SHlo (N * nCls)

The sync-BN data-parallel ResNet-34 forward graph, over the replica family. T2's resnet34FwdGraphB_full with every BatchNorm a bnSyncSiteLA over all R replicas; block prefixes and collective tags are ResNet34RenderB's.

Equations

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

Instances For

theorem Proofs.StableHLO.resnet34FwdGraphSync_full_shard (R : ℕ) (hR : 0 < R) (N : ℕ) (hN : 0 < N) (epsStr : String) {nCls : ℕ} (w : R34BWeights nCls) (e : Fin R → SHlo (N * (3 * (2 * (2 * 56)) * (2 * (2 * 56))))) (X : Vec (R * N * (3 * (2 * (2 * 56)) * (2 * (2 * 56))))) (he : ∀ (r : Fin R), den (e r) = batchShard R N (3 * (2 * (2 * 56)) * (2 * (2 * 56))) X r) (r : Fin R) :

den (resnet34FwdGraphSync_full R hR N epsStr w e r) = batchShard R N nCls (resnet34ForwardB_full (R * N) w X) r

⭐⭐ T2 at synchronised BatchNorm: replica r's forward IS shard r of the global-batch forward. Given that the replicas' inputs are the shards of one batch X of R·N examples, the sync-BN graph on replica r denotes batchShard r of resnet34ForwardB_full (R * N) w X — the committed batch-BN forward, at the global batch. One block lemma per stage, the shard hypothesis threaded from each into the next.