Model dim. Derived so it is DEFEQ to heads * hd — see the note on VitDims.

Tokens including CLS.

ViT-Tiny: D = 192 = 3 × 64, MLP 768. The default everywhere.

ViT-Small: D = 384 = 6 × 64, MLP 1536. Same depth and same patch grid as Tiny — S widens only, which is why it needs no new proof and no new block chain.

ViT-Base (DeiT-B): D = 768 = 12 × 64, MLP 3072. Still depth 12 and still 16×16 patches.

Total encoder blocks = the reference's totalDrop, i.e. the ramp DENOMINATOR is this minus 1.

Two sites per block — the attention branch and the MLP branch, dropping independently.

The mask-input ordinal of block i's branch br (br = 0 attention, 1 MLP). The single source for the numbering: the forward walks blocks upward, the backward downward, and both call this. Distinct from the RAMP index, which is i for BOTH branches.

The ramp index of a site ordinal — the inverse direction, used by the driver's keep table.

The %dp<i>: tensor<Bxf32> inputs an SD ViT render appends to its signature — one per SITE, in ordinal order, which is the order the driver's dropScales writes them into the blob.

Per-block func-arg signature (committed forward order).

ViT-Tiny depth-12 forward rendered ENTIRELY from the verified AST. Every line is pretty of a verified SHlo node; den(graph) = vitForwardKV by vitFwdGraphKMHV_faithful (at depth 12). The output is the [BS,10] logits.

The whole-net backward traversal, SHARED by the SGD and AdamW renders. Returns the emitted code and, in func-arg order, one SSA per parameter — the updated param at adam := false, the un-fused gradient at adam := true. One traversal, two tails: a second copy of the depth-12 backward would be the double-writer disease one level down.

The 200 parameter (name, shape) pairs in func-arg order — the single source for the argument signature, the return types, and (in the AdamW render) the %<nm>m/%<nm>v moment slots.

nClasses is a parameter here, which the matched pair with jax/MainVitImagenet.lean (a 1000-class ViT-Tiny) needs.

V is TRAILING and defaulted to ViT-Tiny, for the same reason vbB is: every existing call site is untouched and every committed artifact re-renders byte-identically.

Does this parameter get weight decay? The renderer's half of the timm rule.

It keys the positional embedding by NAME where the reference keys it by SHAPE (p.shape == _WD_POS_SHAPE), and that is deliberate rather than a transcription slip: the reference walks an unnamed pytree and has nothing else to key on, while a shape test here would also exclude any other param that happened to be 197×192. The name is the more precise identifier when you have one. #guards below pin the resulting counts against the reference's own, which is what stops the two readings drifting.

The rule reads the RANK, and rank is the one thing that survives the layout difference between the two sides — the render carries Wfc1 as [192,768] where the reference has (768,192). Both are 2-D, so both decay.

The decayed / excluded split, as the renderer computes it.

ViT-Tiny depth-12 SGD train step rendered ENTIRELY from the verified AST. Forward (vitFwd12) → softmax-CE cotangent (softmax(logits) − onehot, the lossCotGraph form) → head-dense back (dotOut + weightSgd/biasSgd) → clsPadF → final-LN back (vlnBack) → 12× vBlockBack (reversed, cotangent threaded) → patch-embed back (patchEmbedWeightSgd/patchEmbedBiasSgd

• clsSliceF→denseBiasSgdB for cls + posEmbedSgd for pos). Returns the 200 SGD-updated params in func-arg order. lrStr is the mean-loss-equiv literal (base/BS); cotangent has NO /B (folded into lr). The traversal itself is vitBackAll false, shared with the AdamW render.

The driver's variant slug for a (per-device batch, replica count, EMA) triple: the artifact is verified_mlir/vit_<variant>_train_step.mlir, the entry point is @vit_<variant>_train_step and LEAN_MLIR_VARIANT=<variant> selects it at run time.

This is the ViT peer of cnxAdamVariant / r34AdamVariant / mnv2AdamVariant, and unlike theirs it is documentation plus a drift guard rather than the name's producer — vitAdamTrainStepText takes funcName explicitly and the #eval paths must stay string literals for regen_verified_mlir.sh's writer audit to see them. So the #guards at the bottom of this file are what tie the literals to this function; the contract is checked at lake build rather than merely described.

ViT's spelling breaks the "the number is the per-device batch" convention at 4 replicas (adamdp32x4, adamdp128x4) and that is deliberate — vit_adamdp_train_step.mlir is a COMMITTED 2-replica artifact at bs32, so a 4-replica render reusing adamdp would give one path two writers computing different graphs. Encoded here so the exception cannot be forgotten.

The ema marker LEADS. trainAdamSched keys its 4-region [θ|m|v|ema] blob off variant.startsWith "ema", so a trailing marker would silently select the 3-region layout for a 4-region graph — every parameter misaligned. On EfficientNet the RMSProp+EMA variant is emarms, which does not start with "rms", so a prefix test would initialise the mean-square to 0. ViT is AdamW-only, so there is no second axis here; if one is ever added, make both predicates substring tests first.

ViT-Tiny depth-12 AdamW train step, rendered from the verified AST. The certified peer of the hand-written ViTRender.vitTrainStepModuleAdamSched.

Same backward as vit_train_step (vitBackAll, one traversal) but taking the un-fused gradients, each fed to the proven AdamW triple. The cotangent is the LABEL-SMOOTHED one with an explicit ÷B, matching the AdamW recipe — the SGD render folds the mean into lr and does no smoothing, so the two are different functions and this parameter is not optional.

Interface: 605 in (%x, 200 θ, 200 m, 200 v, %lr/%bc1/%bc2, %onehot) / 603 out (200 θ', 200 m', 200 v', %loss/%bc1/%bc2), the packed [θ|m|v] layout the AdamW trainer reads. Hand-written text besides the signature and constants: the report-only %loss and the %bc1/%bc2 passthroughs.

At ema := true the blob gains a fourth region and the scalar tail goes 3 → 5, so the interface becomes 807 in / 805 out = 605/603 + 200 (the shadow) + 2 (%emad/%oemad). ema is LAST in this signature on purpose: inserted mid-list it would capture an existing positional argument at every call site.