Documentation

LeanMlir.Proofs.Codegen.MobileNetV4RenderB

MobileNetV4 — the Universal Inverted Bottleneck render (planning/archive/mnv4_verified.md phase 3) #

The UIB block, batch-BN, at N := B:

  optional pre-DW (k×k)  → BN → relu      -- takes the block's stride
  expand 1×1 (ic → mid)  → BN → relu
  optional post-DW (k×k) → BN → relu      -- at stride 1 if a pre-DW already consumed it
  project 1×1 (mid → oc) → BN             -- NO activation
  + skip                                  -- iff stride = 1 ∧ ic = oc; NO post-add activation

mid = ic * expand, every conv BN-followed and therefore bias-freeVLayer.toSpecs and the baseline's Layer.nParams both assume that, and uib-layout-tie pins them to each other. ⚠ 8,447,322 scalars in 233 slots at nClasses = 10 (9,715,512 at 1000) over Conv-M's 21-row table — counted off mnv4_fwd.mlir's own signature, not off a docstring. The 3,737,088 / 14-block figure this line carried until 2026-09-07 was Conv-S's, left behind by ed5a797's table swap.

k = 0 omits that depthwise, which is how one block expresses MNv4's four families: ExtraDW (both), IB/MBConv (post only), ConvNeXt-like (pre only), FFN (neither). Those are ifs here, not separate functions, because omitting a shape-preserving op does not change any type.

Phase 0 found this needs no new op and no new position (planning/archive/mnv4_verified.md §2): the depthwise VJP is kernel-general (cnx_render_dw7*_certified; the descriptor carries kH kW), and a leading depthwise already exists — MobileNetV2RenderB.lean:196's t = 1 inverted residual emits .depthwise (c := ic) straight onto the block input. What is new is the composition: ExtraDW puts a depthwise on both sides of the pointwise expand.

Why three functions and not one #

.depthwise : BatchableOp (c*h*w) (c*h*w) and .depthwiseStrided : BatchableOp (c*(2h)*(2w)) (c*h*w) have different INPUT types, so a stride-polymorphic block cannot typecheck — the same reason MobileNetV2RenderB splits irFwdStridedB from irFwdSkipB. The stride-2 case splits again by which depthwise consumes the stride, because that decides the spatial size the expand runs at. Read off the Conv-M table (jax/MainMobilenetV4.lean), which lands cleanly:

strideic vs ocfunction
18 blocks1ic = ocuibFwdSkipB
3 blocks (1, 3, 11)2icoc, preDWk > 0uibFwdPreStridedB
0 blocks2icoc, preDWk = 0uibFwdPostStridedB

⚠⚠ Conv-M has NO post-strided block (Conv-S had one), so that third arm is certified and unexercised — a green corpus is not coverage of it. The split was 11 / 2 / 1 here until 2026-09-07; the #guards in Proofs/Nets/MobileNet/MobileNetV4BackB0.lean have said 18 / 3 / 0 since 2026-08-14.

⚠⚠ ACTIVATION IS PLAIN relu, NOT relu6. MobileNetV2's blocks use relu6 and this file sits next to that renderer, so the wrong one is one keystroke away. Read off the baseline emitter (MlirCodegen.lean:6357, "Plain ReLU throughout").

A pre/post-DW swap is invisible to every count. Same k, same channels ⇒ same parameter shapes, so uib-layout-tie passes on a renderer that swaps them, and so does any arity or op-count audit. At stride 1 it is invisible to the TYPES too, since both positions are shape-preserving — which is why the four families are ifs that the compiler cannot check. Only a forward tie against the reference on shared weights pins the order. Same class as R50's stride-on-the-3×3.

⚠ The baseline drops the stride entirely for a stride-2 FFN block (no depthwise to carry it, MlirCodegen.lean:6364). No such block exists in the table; this file has no function for that shape, so the case is absent rather than silently wrong.

One row of the MobileNetV4-Conv-M block table. h is the block's OUTPUT spatial size, so a stride2 block reads its input at 2h.

  • p : String

    parameter-name prefix: "1""21" (Conv-M; Conv-S ran to "14").

  • ic :
  • oc :
  • expand :
  • preDWk :

    pre-depthwise kernel, 0 = absent.

  • postDWk :

    post-depthwise kernel, 0 = absent.

  • h :
  • stride2 : Bool
Instances For
    Equations
    • One or more equations did not get rendered due to their size.
    Instances For

      THE BLOCK TABLE — transcribed once, from jax/MainMobilenetV4.lean.

      ⭐⭐ Everything downstream folds over this list: the parameter signature, the BN stat slots, the forward chain, the backward chain and the running-statistic recomputes. Before it existed the same rows were hand-written FOUR times, and §3/§7.2's whole point is that a divergence between two such readings is invisible — same ops, same channel counts, same types, different net. One table means a dispatch error is a typo in one place rather than a mismatch nothing checks.

      Families in order (Conv-M): ExtraDW ×7, ConvNeXt, FFN, ConvNeXt, ExtraDW ×4, FFN, ConvNeXt, ExtraDW ×2, FFN ×2, ConvNeXt — 13 ExtraDW / 4 ConvNeXt / 4 FFN, and no IB at all, where Conv-S used three. Spatial ladder 56 → 28 → 14 → 7.

      ⚠ Verified against timm 1.0.28 (mobilenetv4_conv_medium, walking model.blocks[1:4]): all 21 rows agree on (ic, oc, expand, preDWk, postDWk, h, stride2). The #guards in Proofs/Nets/MobileNet/MobileNetV4BackB0.lean pin that reading; they are derived from timm rather than re-read off this table, or they would gate nothing.

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

        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.

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

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

          Equations
          Instances For

            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.

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

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

              Equations
              Instances For

                Saved forward SSA names the UIB backward + gradient passes reference. The optional fields are "" when their depthwise is absent (k = 0), which is how the backward reads off which family it is looking at without re-deriving it from the kernel sizes.

                Instances For

                  Everything the whole-net render needs out of one forward traversal: the emitted code, the logits, and every saved activation the backward reads. inputs is each block's INPUT SSA name (inputs[i] is what block i consumed), which the weight gradients contract against.

                  Instances For
                    def Proofs.StableHLO.mnv4FwdChainB (B nClasses : ) (epsStr : String) (mode : BnMode := BnMode.train) (bf16 : Bool := false) (replicas : := 1) (sync : Bool := false) :

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

                    Transcribed 1:1 from jax/MainMobilenetV4.lean, which is the faithful Conv-M table as of 2026-08-14 (historical/RESULTS.md's 84.58% belongs to the SUPERSEDED Conv-S table). Spatial ladder:

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

                    Block dispatch is forced by the table and checked by the types: the three stride-2 blocks (1, 3, 11) are icoc and split by which depthwise carries the stride — all three are PRE-strided in Conv-M; the eighteen stride-1 blocks are all ic = oc, hence all skip. 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. ⚠ This paragraph listed Conv-S's fourteen families until 2026-09-07.

                    Activations, and they are not uniform — each read off the emitter that produced the number, not assumed:

                    • stem and head .convBnrelu (MlirCodegen.lean:5852, emitConvBnTrain … useRelu := true)
                    • the fused stage → swish (jax/Jax/Codegen.lean:1031) — a deliberate paper deviation, see fusedMbConvFwdStridedB
                    • every UIB block → relu (MlirCodegen.lean:6357, "Plain ReLU throughout")

                    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.

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

                      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.

                      Equations
                      Instances For
                        def Proofs.StableHLO.mnv4FwdFaithfulV (B nClasses : ) (epsStr : String) (slug : String := "mnv4") (vSuffix : String := "") :

                        @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 (§2l step B).

                        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.

                        Equations
                        • One or more equations did not get rendered due to their size.
                        Instances For
                          def Proofs.StableHLO.mnv4FwdEvalFaithfulV (B nClasses : ) (epsStr : String) (slug : String := "mnv4") (vSuffix : String := "") :

                          @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.

                          Equations
                          • One or more equations did not get rendered due to their size.
                          Instances For
                            def Proofs.StableHLO.mnv4AdamVariant (B replicas : ) (bf16 : Bool := false) :

                            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.

                            Equations
                            Instances For
                              def Proofs.StableHLO.mobilenetv4AdamTrainStepFaithfulB (B nClasses : ) (epsStr : String) (replicas : := 1) (slug : String := "mnv4") (bf16 forceSync : Bool := false) :

                              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, which is the only reading that cannot go stale (this line said 583 / 581 at Conv-S's 158 slots until 2026-09-07). 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: stem 3×3/s2 XLA-SAME (3→32, 224→112) → fused MBConv (32→48, 112→56, swish) → the 21 UIB blocks (three stride-2 downsamples, eighteen identity skips; ExtraDW, ConvNeXt and FFN — no IB) → TWO 1×1 conv-BN-relu head stages (256→960, then 960→1280) → GAP → 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.

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