ch10 ViT — verified-faithful StableHLO render fragments (shared library)

Hand-rendered batched StableHLO string fragments for the Vision Transformer (ch7/ch8/ch9 style), each line what the proven-faithful emitter produces — the matmuls are dot_general (proven dense), the row-softmax is the V1 op pattern (softmaxRowF/softmaxRowBack, plain exp/sum), GELU is the ch9 geluF tanh approximation, LayerNorm is layerNormForward (per-token over D, γ=1/β=0) ∘ a per-channel [D] affine (ConvNeXt layerScale + a [D] bias). NOT a single den(trainStep) theorem — faithful PER-OP, validated by the Lean gradchecks (TestSDPA/TestMHSA/TestViTBlock) and by training.

Every fragment is prefix-parameterized (p) so it can be instantiated many times (e.g. 12 distinct blocks) without SSA collisions. A forward fragment KEEPS the intermediate SSA values its backward needs; the matching backward fragment reuses them (after a forward recompute), exactly like ch8's seFwd/seBack. All fragments assume a %sc (f32 0) constant is in scope (the reduce init).

def ViTRender.lnFwd (p x g b : String) (bb n d : ℕ) (eps : String) : String

LayerNorm forward, prefix p, over [b,n,d] (normalize each token over the d axis [2], then per-channel affine). KEEPS %{p}xhat [b,n,d], %{p}istdb [b,n,d], %{p}nf [b,n] for the backward. Result %{p}y.

Equations

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

def ViTRender.lnBack (p g dy : String) (bb n d : ℕ) : String

LayerNorm backward, prefix p. Reuses %{p}xhat/%{p}istdb/%{p}nf from a preceding lnFwd p. dy is the output cotangent, g the γ [d]. Produces %{p}dx [b,n,d], %{p}dg [d], %{p}db [d]. Affine back: dx̂ = dy⊙γ, dγ = Σ dy⊙x̂, dβ = Σ dy. Normalize back (γ=1): dx = istd·(dx̂ − mean_d(dx̂) − x̂·mean_d(dx̂⊙x̂)) (= bn_grad_input with γ=1). (No x arg — the LN input is not needed; x̂/istd come from the fwd recompute.)

Equations

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

def ViTRender.geluActF (p inp : String) (bb n m : ℕ) : String

GELU forward (geluF lines) over [b,n,m], prefix p, input SSA inp. Result %{p}a.

Equations

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

def ViTRender.geluActBack (p xpre dy : String) (bb n m : ℕ) : String

GELU backward (geluBack lines) over [b,n,m], prefix p. Recomputes gelu'(x) from the saved pre-activation xpre; dy is the cotangent. Result %{p}dx.

Equations

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

def ViTRender.mlpFwd (p x Wfc1 bfc1 Wfc2 bfc2 : String) (bb n d m : ℕ) : String

MLP forward, prefix p. KEEPS %{p}h1 [b,n,m] (pre-gelu) and %{p}ga [b,n,m] (post-gelu) for the backward. Result %{p}y [b,n,d].

Equations

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

def ViTRender.mlpBack (p x Wfc1 Wfc2 dy : String) (bb n d m : ℕ) : String

MLP backward, prefix p. Reuses %{p}h1/%{p}ga from mlpFwd p. dy [b,n,d]. Produces %{p}dx [b,n,d] + %{p}dWfc1 [d,m], %{p}dbfc1 [m], %{p}dWfc2 [m,d], %{p}dbfc2 [d].

Equations

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

def ViTRender.mhsaFwd (p x Wq bq Wk bk Wv bv Wo bo : String) (b n d h dh : ℕ) (scale : String) : String

MHSA forward, prefix p. Saves %{p}W [b,h,n,n], %{p}Qh/%{p}Kh/ %{p}Vh [b,h,dh], %{p}P [b,n,d] for the backward. Result %{p}O.

Equations

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

def ViTRender.mhsaBack (p x Wq Wk Wv Wo dO : String) (b n d h dh : ℕ) (scale : String) : String

MHSA backward, prefix p. Reuses the mhsaFwd p saves. Produces %{p}dx [b,n,d] + 8 param grads %{p}dWQ %{p}dbQ %{p}dWK %{p}dbK %{p}dWV %{p}dbV %{p}dWo %{p}dbo. (Validated in tests/TestMHSA.lean, rel err 2e-3.)

Equations

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

structure ViTRender.BlockParams : Type

A block's 16 parameter SSA names, in layout order: LN1 γ/β [d]; MHSA Wq/bq/Wk/bk/Wv/bv/Wo/bo [d,d]/[d]; LN2 γ/β [d]; MLP Wfc1/bfc1/Wfc2/bfc2 [d,m]/[m]/[m,d]/[d].

• g1 : String
• b1 : String
• Wq : String
• bq : String
• Wk : String
• bk : String
• Wv : String
• bv : String
• Wo : String
• bo : String
• g2 : String
• b2 : String
• Wfc1 : String
• bfc1 : String
• Wfc2 : String
• bfc2 : String

def ViTRender.blockFwd (p x : String) (bp : BlockParams) (b n d m h dh : ℕ) (eps scale : String) : String

Transformer block forward, prefix p. x [b,n,d]. Result %{p}out. Keeps %{p}r1 (the first residual) + all sub-fragment saves. The sub-prefixes are {p}1 (LN1), {p}m (MHSA), {p}2 (LN2), {p}p (MLP).

Equations

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

def ViTRender.blockBack (p dOut : String) (bp : BlockParams) (b n d m h dh : ℕ) (scale : String) : String

Transformer block backward, prefix p. Reuses blockFwd p saves. dOut [b,n,d] is the block-output cotangent. Produces %{p}dx [b,n,d] and the 16 param grads (see blockGradNames).

Equations

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

def ViTRender.blockGradNames (p : String) : List String

The 16 block param-grad SSA names produced by blockBack p, in the same (layout) order as BlockParams: g1,b1, Wq,bq,Wk,bk,Wv,bv,Wo,bo, g2,b2, Wfc1,bfc1,Wfc2,bfc2.

Equations

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

def ViTRender.patchEmbedFwd (p x w bias : String) (b ic d ph pw s : ℕ) : String

Patch-embed conv forward, prefix p: non-overlapping s×s/stride-s conv [b,ic,s·ph,s·pw]·[d,ic,s,s] + bias → [b,d,ph,pw], then flatten to tokens [b, ph·pw, d] (transpose [0,2,3,1] + reshape). Result %{p}tok.

Equations

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

def ViTRender.patchEmbedBack (p x dtok : String) (b ic d ph pw s : ℕ) : String

Patch-embed backward (weight + bias grads; NO input grad — first layer), prefix p. dtok [b,ph·pw,d] is the token-grad. Reshapes/transposes back to [b,d,ph,pw], then: bias grad reduce[0,2,3]→[d]; weight grad = dilate dy interior s-1 (no high → s·ph-(s-1)), valid conv → [d,ic,s,s]. Produces %{p}dw [d,ic,s,s], %{p}db [d].

Equations

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

def ViTRender.clsPosFwd (p tok cls pos : String) (b n0 d : ℕ) : String

CLS-token + positional-embed forward, prefix p: prepend a learned [d] CLS at row 0 of the n0 patch tokens → [b,n0+1,d], then add a learned [n0+1,d] positional embedding (broadcast over batch). Result %{p}z [b,n0+1,d].

Equations

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

def ViTRender.clsPosBack (p dz : String) (b n0 d : ℕ) : String

CLS + pos backward, prefix p. dz [b,n0+1,d]. Produces %{p}dtok [b,n0,d] (patch-token grad), %{p}dcls [d] (Σ over batch of row 0), %{p}dpos [n0+1,d] (Σ over batch).

Equations

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

def ViTRender.headFwd (p z Wc bc : String) (b n d nc : ℕ) : String

Classifier head forward, prefix p: take the CLS token (row 0 of [b,n,d]) → [b,d], dense [d,nc] + bias → [b,nc] logits. Result %{p}logits.

Equations

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

def ViTRender.headBack (p Wc dlog : String) (b n d nc : ℕ) : String

Classifier head backward, prefix p. Reuses %{p}clsv from headFwd. dlog [b,nc]. Produces %{p}dz [b,n,d] (cotangent scattered into row 0, zero elsewhere), %{p}dWc [d,nc], %{p}dbc [nc].

Equations

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

structure ViTRender.ViTConfig : Type

ViT hyperparameters (one config). n0 = ph·pw patches, n = n0+1 tokens.

• b : ℕ
• ic : ℕ
• d : ℕ
• ph : ℕ
• pw : ℕ
• s : ℕ
• m : ℕ
• h : ℕ
• dh : ℕ
• nc : ℕ
• eps : String
• scale : String

def ViTRender.vitFwd (p x wConv bConv cls pos gF bF Wc bc : String) (blocks : List BlockParams) (cfg : ViTConfig) : String

Whole-ViT forward, prefix p: x [b,ic,s·ph,s·pw] → logits [b,nc]. Result %{p}hdlogits. Keeps every sub-fragment save for the backward.

Equations

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

def ViTRender.vitBack (p dlog x _wConv Wc gF : String) (blocks : List BlockParams) (cfg : ViTConfig) : String

Whole-ViT backward, prefix p. Reuses vitFwd p saves. dlog [b,nc] is the logits cotangent. x is the (fixed) input image (for the patch weight-grad). Produces all param grads (see vitGradNames); NO image grad (first layer).

Equations

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

def ViTRender.vitParamNames (blocks : List BlockParams) : List String

Param SSA names in canonical (layout) order: wConv, bConv, cls, pos, then each block's 16 (via BlockParams fields), then gF, bF, Wc, bc.

Equations

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

def ViTRender.vitGradNames (p : String) (blocks : List BlockParams) : List String

Param grad SSA names produced by vitBack p, SAME order as vitParamNames.

Equations

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

def ViTRender.vitParamDims (blocks : List BlockParams) (cfg : ViTConfig) : List (List ℕ)

Param dims in canonical order (matches vitParamNames).

Equations

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

def ViTRender.vitParamSig (blocks : List BlockParams) (cfg : ViTConfig) : String

The param func signature %nm: tensor<…> (canonical order).

Equations

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

def ViTRender.vitFwdModule (cfg : ViTConfig) (blocks : List BlockParams) : String

@vit_fwd(%x flat, params…) → logits [b,nc] — image→logits (for eval). The flat %x [b, ic·H·W] is reshaped to [b,ic,H,W] then run through vitFwd.

Equations

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

def ViTRender.vitTrainStepModule (cfg : ViTConfig) (blocks : List BlockParams) (lr : String) : String

@vit_train_step(%x flat, params…, %onehot) → updated params — one mean-loss-SGD step. Forward → softmax-CE cotangent dy=(softmax(logits)−onehot)/b → vitBack → per-param SGD θ ← θ − lr·dθ (baked in), returns the updated param list.

Equations

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

def ViTRender.vitTinyBlocks (depth : ℕ) : List BlockParams

Production ViT-Tiny config @ Imagenette 224² (matches ViTLayout): depth-k blocks with distinct per-block param names %<field>_<i>.

Equations

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

def ViTRender.vitTinyConfig (b _depth : ℕ) : ViTConfig

Equations

• ViTRender.vitTinyConfig b _depth = { b := b, ic := 3, d := 192, ph := 14, pw := 14, s := 16, m := 768, h := 3, dh := 64, nc := 10, eps := "1.0e-5", scale := "0.125" }

def ViTRender.emitAdamV (θ g m v : String) (ds : List ℕ) (t : String) : String × String × String × String

AdamW update for one parameter. Emits the m'/v'/θ' block at shape ds (tag t keeps SSA names distinct), reading scalar args %b1 %ob1 %b2 %ob2 %bc1 %bc2 %lr %eps %wd. Returns (ir, θ'SSA, m'SSA, v'SSA). Op-for-op the coordinate formula Proofs.adamWParam: θ' = θ − lr·((β₁m+(1−β₁)g)/bc₁)/(√((β₂v+(1−β₂)g²)/bc₂)+ε) − (wd·lr)·θ.

Equations

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

def ViTRender.emitGradAllReduce (g : String) (ds : List ℕ) (t : String) (replicas : ℕ) : String × String

Cross-replica gradient mean: all_reduce(add) over replicas devices, then divide by replicas. Emitted BEFORE the optimizer consumes the gradient, so every replica applies an identical update and the parameter copies stay in lockstep with no host round trip.

⭐ Since 4d piece 2 (2026-09-07) every Proofs/Codegen render calls Proofs.StableHLO.prettyAllReduceMean instead — pretty of the allReduceMeanF AST node, whose token emit (allReduceMeanText) is this body verbatim, so the artifacts did not move and the collective is inside the faithfulness theorems. This function stays for the hand-written emitAdamV path below.

At replicas = 1 this emits nothing and returns the gradient unchanged, so single-device renders stay byte-identical.

Syntax validated end to end by ffi/test_pjrt_allreduce.c. Note the absence of use_global_device_ids: setting it requires a positive channel_id, and for a plain cross-replica reduce it is not wanted (planning/archive/xla_pjrt_ladder.md §11).

Equations

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

def ViTRender.emitAdamVDP (θ g m v : String) (ds : List ℕ) (t : String) (replicas : ℕ := 1) : String × String × String × String

emitAdamV with the gradient first averaged across replicas devices — the data-parallel AdamW update. replicas = 1 is exactly emitAdamV.

The proofs are untouched: each replica evaluates the same tied graph at the batch size it was rendered for, and the collective averages gradients of that function over disjoint equal batches (planning/archive/xla_pjrt_ladder.md §10.4). Prefer SCALING the global batch over splitting it — that keeps BatchNorm's group size, and therefore the tie, unchanged.

Equations

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

def ViTRender.vitTrainStepModuleAdam (cfg : ViTConfig) (blocks : List BlockParams) : String

@vit_train_step_adam — the SGD train step's optimizer swapped for AdamW. Same forward/backward/softmax-CE cotangent as vitTrainStepModule; the per- param SGD θ−lr·dθ is replaced by emitAdamV (so the func also takes the per-param moments %<nm>m/%<nm>v and the scalar Adam hyperparameters). Returns the updated parameters (moment outputs elided for the smoke; a full step would also return %admn/%advn).

Equations

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

def ViTRender.vitTrainStepModuleAdamPacked (cfg : ViTConfig) (blocks : List BlockParams) (lr β1 ob1 β2 ob2 eps wd : String) : String

Packed AdamW train step for the FFI driver. Hyperparameters are baked as constants (so the func takes NO scalar args), and the parameters + both moment buffers thread as a single [θ|m|v] blob: arg order (x, θ×k, m×k, v×k, onehot) and return (θ'×k, m'×k, v'×k). This matches iree_ffi_train_step_generic's (x, params, y) → params' contract with n_params = 3k (the moments ride in the params blob — no .so change). Bias correction is omitted (bc₁=bc₂=1); a later rung host-passes the per-step 1−βᵗ. Optimizer = Proofs.adamWParam.

Equations

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

def ViTRender.vitTrainStepModuleAdamSched (cfg : ViTConfig) (blocks : List BlockParams) (β1 ob1 β2 ob2 eps wd : String) (ls : Float) : String

Scheduled AdamW train step (Phase 2): like …AdamPacked, but lr/bc₁/bc₂ arrive as runtime rank-0 scalar params (smuggled in the packed blob's tail — the FFI takes no scalar slot) so the host can drive cosine+warmup and the per-step bias correction 1−βᵗ. They are returned UNCHANGED (passthrough) so #outputs = #inputs = 3k+3, preserving the generic FFI's invariant. Arg order (x, θ×k, m×k, v×k, lr, bc₁, bc₂, onehot); only β₁,β₂,ε,wd stay baked.

Equations

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