The sync forward is POINTWISE once its statistics are fixed.

Cell t becomes γ_c·(x_t − μ_c)·(m2_c − μ_c² + ε)^(−1/2) + β_c with c its own channel — no other cell is read. Cells mix ONLY when the statistics are computed, and under sync-BN that computation has been hoisted out into the collective.

This is what makes the shard argument (bnSyncTensor4_batchShard) pure index bookkeeping: a pointwise map commutes with any reindexing that preserves the channel, and sharding the batch does.

Sharding the batch does not move a cell's CHANNEL. The one fact P1 needs about the layout: the batch axis is outside the channel axis in [N,C,H,W], so cutting the batch leaves every cell in the channel it was already in.

The global channel row, restricted to replica r's block, IS that replica's own channel row. The layout fact that connects a batch shard to the [C, N·H·W] world bnPerChannelFlat reduces in — i.e. the one step the bnchwFwd relabel was hiding.

P1a — the sync forward COMMUTES WITH SHARDING, for any statistics at all.

shard_r ∘ (sync BN at μ, m2) = (sync BN at μ, m2) ∘ shard_r. No hypothesis on μ/m2: once the statistics are fixed the map is pointwise (bnSyncTensor4_apply) and sharding preserves each cell's channel (bnchwChan_batchShard), so there is nothing to prove about BatchNorm here — only about indices.

This is the half of P1 that carries no mathematics, and separating it is what leaves the real content in one place: whether the handed-in statistics ARE the global ones, which is bnMean_shard/bnMeanSq_shard, i.e. exactly what allReduceMeanF computes.

P1b — the GLOBAL per-channel mean is the mean of the replicas' per-channel means. bnMean_shard transported along the row shard. This is precisely what syncStats's first collective — allReduceMeanF over the replicas' bnBatchMeanB — computes.

…and so is the global per-channel SECOND MOMENT. Note: There is no such statement for the variance alone — the mean of the shards' variances is not the variance of the union — which is why the exchange carries Chan's corrected variance (bnVar_row_shard_chan) rather than a plain σ²_r.

Chan's parallel variance on the channel rows: the global channel variance is the replica mean of each replica's own two-pass variance plus its mean's squared offset from the global mean. This is what the second collective of a sync-BN forward carries (StableHLO.bnBatchVarAtB), and it is why no consumer ever forms E[x²] − μ².

P2b's workhorse: a mean over the global batch of ANY pointwise function of two rows is the mean of the replicas' means of the same. Both dy-reductions the sync backward needs have this shape — mdy reads only dy, mdyx reads x and dy together — so one lemma covers both, and bnMean_row_shard is its one-row special case.

The global batch-spatial count is nonzero when a replica's is: (R·N)·(h·w) = R·(N·(h·w)).

P1 — SYNC-BN ON REPLICA r IS THE SHARD-r BLOCK OF THE GLOBAL-BATCH BN.

Handed the GLOBAL statistics — (1/R)·Σ_r of each replica's own bnMean and bnMeanSq, which is what syncStats denotes once its σ² is read back as m2 = σ² + μ² — replica r's sync forward on its own shard equals batchShard r of bnBatchTensor4 run on the whole R·N batch.

The spec does not move: the right-hand side is the EXISTING bnBatchTensor4, at N := R·N. Nothing new is being specified; the render is being shown to hit a target the tier already names.

It is P1a (pointwise, so sharding commutes — no mathematics) composed with P1b (the statistics really are the all-reduced ones — all the mathematics) through bnSyncTensor4_at_own_stats.

A channel's γ gradient over the global batch is the SUM of the replicas' γ gradients at the same handed-in statistics — a sum, not a mean, because this is a parameter gradient: the parameter collective's 1/R is what turns it into the global-batch mean. The row split bnchwFwd_row_batchShard, under Σ instead of bnMean.

…and so is the β gradient, which reads no statistic at all: Σ dy over the global channel row is the sum of the shards' Σ dy.

The sync backward is POINTWISE too, with the channel read off the index — it reads only x idx and dy idx. Both its reductions were hoisted into the collective.

The [N,C,H,W] lift: cell t of the sync backward depends only on x t, dy t and its own channel's four statistics.

P2a — the sync backward COMMUTES WITH SHARDING, for any statistics at all. The backward twin of bnSyncTensor4_batchShard, and equally free of mathematics: pointwise plus channel-preserving reindex.

The R = 1 backward anchor restated with bnSyncXhat in the mdyx reduction — the form the sync GRAPH produces, since bnSyncDyStatsB builds x̂ from the statistics handed to it rather than from x directly.

P2 — THE SYNC BACKWARD ON REPLICA r IS THE SHARD-r BLOCK OF THE GLOBAL-BATCH INPUT-VJP.

Handed the four all-reduced statistics — each literally (1/R)·Σ_r' of a per-replica quantity, which is what syncStats, then allReduceMeanF of bnSyncDyStatsB, denotes — replica r's sync backward equals batchShard r of bnBatchTensor4GradInput run on the whole R·N batch.

Including the cross-shard terms. mdyx averages x̂·dx̂ with x̂ built from the GLOBAL μ, m2 — not from the shard's own statistics — which is exactly why bnSyncDyStatsB consumes the already-reduced vector instead of recomputing one. That is what makes each replica's output the true shard-r block of the global gradient, and hence the all-reduced PARAMETER gradient exact rather than approximate.

The spec does not move: the right-hand side is the existing bnBatchTensor4GradInput at N := R·N.

Example n of shard r is example (r, n) of the global batch.

batchMap commutes with sharding — the block analogue of batchSlice_batchMap. A per-example lift applied to the global batch, restricted to replica r, is the same lift applied to replica r's shard: so every conv, relu, pool, GAP and dense node in a DP render denotes batchShard r of its batch-R·N value as soon as its input does.

…and so does batchMapAux, the shape of every batched backward that recomputes from a saved per-example activation: example n is handed ITS slice of aux, which sharding respects.

A pointwise map commutes with sharding (scaleB, shiftB, divConstB, …).

A pointwise binary map commutes with sharding (addVB, subB, the relu mask selectPosB against its saved activation, …).

The replica mean of a replica-independent value is that value. bnSyncDyStatsB passes the already-global [μ ‖ m2] through the second collective, and this is why that costs nothing.

The first collective IS the global mean — bnMean_shard on the channel rows.

The second collective IS the global variance — Chan's parallel variance (bnVar_row_shard_chan): each replica's two-pass σ²_r plus (μ_r − μ)², averaged. This is the lemma a one-round E[x²] exchange does not have.

σ² + μ² at the global statistics is the global second moment — how a consumer's den, stated at (μ, m2), reads the packed [μ ‖ σ²].

P1 on the graph, any R. Replica r's bnSyncF, fed by the two-round statistics subgraph over the R replicas' inputs, denotes batchShard r of the batch-R·N bnBatchTensor4 — given that each replica's operand denotes its shard (hx, the chain's induction hypothesis). den_bnSyncF_allReduce_R1 is this at R := 1.

P2 on the graph, any R. Replica r's bnSyncBack, fed by the collective over the replicas' bnSyncDyStatsB (each reading the packed forward statistics), denotes batchShard r of the batch-R·N bnBatchTensor4GradInput — given that each replica's saved activation and its incoming cotangent are its shards of the global ones (hx / hxv, hdy). den_bnSyncBack_allReduce_R1 is this at R := 1.

P2γ on the graph, any R, already all-reduced. The parameter collective over the replicas' bnSyncGammaGradB — each at its shard, its shard's cotangent and the packed global statistics — is 1/R of bnPerChannelGradGamma at N := R·N: the committed γ gradient at the global batch. The BN γ node is the one parameter gradient sync-BN changes, because it is the one that reads x̂.

The handed-back running MEAN under sync-BN is the global batch's own — bnStatsMeanB on the packed statistics denotes what bnBatchMeanB at N := R·N denotes, so the host EMAs the global statistic on every replica.

…and so is the handed-back VARIANCE — the global batch's bnVar, what bnBatchVarB at N := R·N denotes. Note: No per-replica variance is ever averaged on its own: the between-shard spread (μ_r − μ)² rides along, which is what Chan's formula is.

P4 for the Σ_n-shaped gradients, at the conv weight. The collective over the replicas' convWeightGradB, each on its shard at the shard-r block of the global cotangent, is 1/R of the batch-R·N node at that cotangent. den_allReduceMeanF_convWeightGradB (DataParallel.Node) is the same collective with NO relation between the replicas' inputs; this is what it becomes once P1–P3 relate them. Every other Σ_n gradient (denseWeightGradB, denseBiasGradB, the strided / depthwise / bias kinds) closes by the same shard_sum.

P4 at the BN β gradient: Σ dy over each shard's channel row, averaged, is 1/R of the global row's Σ dy. β reads no statistic, so this is the row split alone.

A VJP backward is linear in its cotangent — read off HasVJP.correct. This is the one step that reconciles a DP render's divConstB N with the batch-R·N step's divConstB (R·N): the per-replica cotangent is R · the shard of the global one, so every per-replica gradient is R · its shard contribution, and the collective's 1/R cancels it.