The stochastic-depth sites: the block indices (into mnv4Blocks) of the 18 skip rows, in order. Site k is %dp<k>, and the driver's dropKeeps lists one keep per site in this order. The three strided rows (0, 2, 10) have no residual, so no site, as in the reference's uib_block (_drop_branch only under residual.shape == x.shape).

Block bi's drop site, none for a strided row.

One UIB block's parameters, in VLayer.toSpecs order: pre-DW? → expand → post-DW? → project, each {W, γ, β} and bias-free. The names are exactly what uibFwd*B emits, and the order is exactly what toSpecs lays out — those two facts are what make the signature and the driver's parameter blob describe the same thing.

A depthwise kernel is [c, 1, k, k], not [c, c, k, k]. That is the whole difference between a depthwise and a regular conv at this layer, and it is 3 orders of magnitude of parameters.

The fused stage's parameters, in toSpecs order: the k×k regular conv then the 1×1 project.

The MobileNetV4-Conv-M parameter inputs, in func-arg order: stem (3), the fused stage, the 21 UIB blocks, the TWO head convs, the classifier. Single source for the signature and the return order, the same role r50ShapeList plays for R50.

This list and VLayer.toSpecs are TWO HAND-WRITTEN READINGS of the same layout — the renderer cannot import the spec without inverting the dependency, which is the same two-lists shape as toSpecs == XLayout.specs elsewhere. mnv4-fwd-smoke is the gate that pins them.

The same list as MLIR types. Derived, so the shapes have one definition.

One UIB block's BN stat slots, in the SAME order uibSig lays out its parameters — pre-DW (at ic), expand (at mid), post-DW (at mid), project (at oc), each present exactly when its conv is. The ifs are the same k = 0 dispatch, so a family that drops a depthwise drops its stat slot too and the two lists cannot go out of step.

The fused stage's BN stat slots — the k×k conv's, then the 1×1 project's.

The 154 BN running-statistic slots = 77 BN layers × (μ, var), in forward-traversal order: stem, the fused stage's two, each UIB block's (2–4 depending on family), the head.

A misaligned stat slot is SILENT: the arities still match and the wrong layer's statistics simply flow into the wrong @mnv4_fwd_eval slot. That is why the order here and the order the train step returns them in are both derived from the same block table, and why the eval forward reads them through mnv4Bn's statP rather than an independently-numbered list.

The stat slots as MLIR types. Derived, so the widths have one definition.

One BN site at the batched index — the ONLY place the two BN worlds are chosen between.

statP names the running-stat slot: in .eval it is read (%{statP}mu/%{statP}var), in .train it is the slot the train step hands the batch statistics back in. Both modes call pretty exactly once, so the fresh-name counter advances identically and @mnv4_fwd re-renders byte-identical after this threading — the cheap self-check that it is inert.

This exists so @mnv4_fwd and @mnv4_fwd_eval cannot be different nets: a forward and the train step that trains it can otherwise sit in different BN worlds with every byte-pairing check green. One traversal, one switch, as in ResNet50RenderB.

SYNC-BN: .train is bnFwdSite, which at sync emits the all-reduced composition and returns the packed [μ ‖ σ²] as the third component ("" otherwise, and always at .eval). The collectives are tagged with the γ parameter's name without %.

Stride-1 UIB with the identity skip (ic = oc = c) — 18 of the table's 21 blocks, and three of the four families (Conv-M uses no IB). Everything runs at h×h: both depthwise positions are shape-preserving, so the k = 0 omissions are plain ifs. Block output = addVB (project-BN out) (block input); the bottleneck is LINEAR, no activation after the add.

Stride-2 UIB: the POST-DW carries the stride (timm's dw_mid). An optional pre-DW (dw_start, BN only, stride 1) and the expand run at the INPUT size 2h×2h; the post-DW then downsamples to h×h. No skip — ic ≠ oc. All three Conv-M downsamples (rows 1, 3, 11) are ExtraDW, so they take this form with the pre-DW present; a strided row without a post-DW does not occur in Conv-M and is refused by the #guard below.

Fused inverted bottleneck, stride 2, no SE — MobileNetV4's stage 0 (.fusedMbConv 32 48 4 3 2 1 false) and EfficientNetV2's early stages.

  k×k regular conv ic→mid at stride 2  → BN → relu      -- 2h×2h → h×h
  1×1 project mid→oc                   → BN             -- NO activation

"Fused" means the MBConv expand-1×1 and its depthwise collapse into ONE regular k×k conv, so despite living in a mobile net there is nothing depthwise here. No skip: ic ≠ oc and stride 2.

ReLU, as timm's EdgeResidual in mobilenetv4_conv_medium (scripts/parity/mnv4_timm_parity.py). The JAX fused_mbconv_block is shared with EfficientNetV2 (swish) and takes the activation as an argument.

Block-shape dispatch, in one place. Which of the two forwards a row uses is forced by the row: stride 1 ⇒ the identity skip (ic = oc, pinned by mnv4-fwd-smoke); stride 2 ⇒ the post-DW carries it (timm's rule when both depthwise convs exist). The two cannot be one function — .depthwise and .depthwiseStrided differ in INPUT type.

Stem forward: 3×3/s2 conv (3→32), 224→112 → batch BN → relu, on %x. Symmetric (1,1) padding (.convStridedAt), as timm's conv_stem. XLA-SAME padding (.convStridedXlaAt, (0,1) for a 3×3/s2 at 224) also gives 112×112, so only a forward on shared weights sees the difference (scripts/parity/mnv4_timm_parity.py, scripts/parity/mnv4_forward_tie.py).

Head forward, timm's order: 1×1 (256→960) → BN → relu at 7×7 (Conv-M's cn_r1_k1_s1_c960), GAP(7×7), then conv_head 1×1 (960→1280) → norm_head BN → relu on the POOLED [B, 960, 1, 1], then dense. The second BN therefore normalises over the batch alone (N·1·1 cells per channel). Running conv_head at 7×7 before the pool would be a different function with the same parameter count: batch BN and relu do not commute with pooling.

The MobileNetV4-Conv-M forward chain, batch BN, at N := B, 224² → nClasses.

Transcribed 1:1 from the Conv-M table in jax/MainMobilenetV4.lean. Spatial ladder:

  224 --stem s2--> 112 --fused s2--> 56 --uib s2--> 28 --uib s2--> 14 --uib s2--> 7 --GAP--> 1

Block dispatch follows the table: the three stride-2 blocks (1, 3, 11) are ic ≠ oc and strided at the post-DW (timm's dw_mid); the eighteen stride-1 blocks are all ic = oc, hence all skip. The stride/skip split is checked by the types; which depthwise positions a block has (its family) is an if the types do not check. Families in order after the fused stage: ExtraDW ×7, ConvNeXt, FFN, ConvNeXt, ExtraDW ×4, FFN, ConvNeXt, ExtraDW ×2, FFN ×2, ConvNeXt — 13 / 4 / 4 and no IB at all.

Activations, as timm's mobilenetv4_conv_medium:

• stem and head .convBn → relu (MlirCodegen.emitConvBnTrain … useRelu := true)
• the fused stage → relu (timm's EdgeResidual)
• every UIB block → relu after the expand and the post-DW; the pre-DW (dw_start) and the project are BN only (timm's UniversalInvertedResidual)

Returns the full forward RECORD, not just (code, logits): the train step needs every saved activation the backward reads, and the alternative — a second copy of the chain inside the train step — is the two-readings defect this file exists to avoid.

Every distinct channel width a bias-free conv in this net binds %zb{c} at: the stem, the fused stage's two convs, each block's pre-DW (ic), expand and post-DW (mid) and project (oc), and the head. Derived by hand and pinned by mnv4-fwd-smoke, which fails on an unbound %zb.

@mnv4_fwd — the MobileNetV4-Conv-M forward as one MLIR module.

%x plus mnv4SigList's parameters in VLayer.toSpecs order, logits [B, nClasses]. Every conv is bias-free, so the proven conv ops' bias operands bind to the %zb{c} zero constants the prelude declares — same op, bias = 0, and x + 0.0 is exact.

zeroBiasPrelude must cover every width the chain references or the module has an unbound SSA name. That is a link error at parse time rather than a wrong number, which is the good direction, but mnv4-fwd-smoke checks it anyway so the failure arrives at lake build and not at iree-compile.

@mnv4_fwd_eval — the inference forward, every BN site reading frozen running stats. %x + 233 params + 154 stat inputs = 388 inputs (counted off mnv4_fwd_eval.mlir). This is what the driver scores through.

It is mnv4FwdChainB at .eval — the SAME traversal @mnv4_fwd and the train step use, so its BN order matches mnv4StatSigList by construction rather than by a second reading. Its typed graph is mnv4FwdGraphBFullEval at f = s / 32 (MobileNetV4FullBEval).

Pair a block's uibSig/fusedSig slice with the gradient SSA names its backward produced, in the same order.

ONE source for the names AND the shapes. The peers hand-write the parameter list twice — once in the signature, once in the backward's ps — and rely on eyes to keep them in step. Here the list the func header is built from is literally the list the AdamW tail folds over, so a parameter cannot be one shape in the signature and another in the optimizer. The k = 0 dispatch is then written once, in uibSig, instead of once per family per direction.

A bias-free 1×1 conv + BN, reversed — the pass every UIB/fused backward runs for its project and expand stages: the BN input-VJP (bnBackSite), the conv weight gradient on the BN-back output, γ (bnGammaSite) and β (bnBetaGradB) on the incoming cotangent dyIn, and the conv input-VJP. ic → oc at h×h; wName/gName/tag are the stage's weight, γ and sync-BN tag, xc and st the forward's saved conv output and statistics, actIn the conv's forward input. The caller passes dyIn already masked (selectPosB) where an activation sits between stages. Returns (code, dW, dγ, dβ, dx), the code in pretty-call order, so the SSA numbering is the inlined spelling's.

STRIDE-1 UIB backward + its 6/9/12 un-fused gradients — all four families, ic = oc = c. The skip is linear, so dy reaches the project BN unchanged AND fans back in at the block dx.

The two ifs are the same k = 0 dispatch the forward used. A backward that omits the other depthwise from the one the forward emitted still type-checks (both positions are shape-preserving at stride 1) and still descends — it computes the gradient of a different net. So the forward record f is read for which positions exist (f.qr/f.dn are "" when absent) rather than re-deriving it.

STRIDE-2 UIB backward — the POST-DW carries the stride (uibFwdStridedB). The optional pre-DW and the expand run at the INPUT resolution 2h×2h, so their weight gradients contract at 2h while the project's runs at h. The pre-DW is BN only, so its cotangent is the expand's input-VJP unmasked. No skip (ic ≠ oc).

Fused inverted-bottleneck backward, stride 2 — MobileNetV4's stage 0.

The forward is relu (timm's EdgeResidual), so the mask is selectPosB at the BN output.

Backward dispatch — the SAME ifs as uibFwdDispatch, on the SAME row.

That the two dispatchers read one UibSpec is the point. Differentiating a block as a different family than the forward emitted it type-checks at stride 1 (both depthwise positions are shape-preserving), produces the right shapes, trains and descends — and computes the gradient of a net nobody built. With one table and one row per call there is no second place for the dispatch to be written down differently.

The driver's variant slug for a given (B, replicas): the artifact is verified_mlir/mnv4_<variant>_train_step.mlir, the entry point is @mnv4_<variant>_train_step, and LEAN_MLIR_VARIANT selects it. All three must agree — the shim checks the entry name and refuses a mismatch outright rather than running the wrong graph. B = 32 is deliberately unsuffixed so the Imagenette artifact keeps a stable name.

The variant slug of a recipe render — the JAX reference's recipe renders. opt = none is mnv4AdamVariant, so every committed spelling is unchanged. Markers in the order the driver's predicates read them (VerifiedVariant): ema first (emaOn is a PREFIX test), then the optimizer with its k (accK parses it back out after acc[dp]), the per-replica batch, wx, do, the decay mark, bf16.

MobileNetV4-Conv-M AdamW train step, batch BN, rendered from the verified AST at N := B.

858 inputs (%x, 233 θ, 233 m, 233 v, %lr/%bc1/%bc2, 154 running-stat slots, %onehot) and 856 outputs (233 θ', 233 m', 233 v', %loss/%bc1/%bc2, 154 batch stats) — counted off mnv4_adam_train_step.mlir. Parameter order comes from mnv4ShapeList — through zipPs, which builds each block's gradient list from the very same uibSig slice the signature does — and stat order from mnv4StatShapeList.

Forward, as timm's mobilenetv4_conv_medium: symmetric stem 3×3/s2 (3→32, 224→112) → fused MBConv (32→48, 112→56, ReLU) → the 21 UIB blocks (stride on the post-DW, BN-only pre-DW; ExtraDW, ConvNeXt and FFN — no IB) → 1×1 conv-BN-relu (256→960) → GAP → 1×1 conv_head (960→1280, BN over the batch at 1×1) → dense.

The cotangent is composed from kit ops (softmaxRow → subB → scaleB → addVB → shiftB → divConstB, α = 0.1, K = nClasses), and %loss is report-only and stays outside the AST — the same carve-out resnet34/mobilenetv2 take.