Documentation

LeanMlir.Proofs.Codegen.ConvNeXtRenderB

ConvNeXt at the BATCHED index N := B — forward, backward and AdamW renders #

The ConvNeXt peer of ResNet34RenderB / MobileNetV2RenderB, and the reason it exists is stochastic depth: the drop mask is per-EXAMPLE, and in the per-example-indexed render (ConvNeXtRender) a node denotes ONE example — pretty B lifts it across the batch, so the node cannot see j. dropPathB needs its operand at index B·n, which is what this file's chain produces.

The trap this closes is that the wrong thing TYPECHECKS. pretty B already emits tensor<B×n>, so a broadcast_in_dim %mask, dims = [0] + multiply against a per-example node compiles, trains and descends — with no faithful den behind it. Every node below is instead a batchOp/*B form whose den is batchMap N (…) or batchMapAux N (…), i.e. honest about which index is the batch.

What this file writes — every ConvNeXt artifact but one. The forwards (drop-free and stochastic-depth), and the AdamW/EMA train steps through ConvNeXtRender.convNextAdamTrainStepText with this file's traversal, at ConvNeXt-T/S/B, Imagenette and ImageNet, f32 and bf16. Only the SGD-inline convnext_train_step.mlir is written by ConvNeXtRender: this traversal has no fused-SGD arm, and ConvNeXtStepTie's 182-parameter tie is stated at those bytes. The Proofs tier for this chain is Nets/ConvNeXt/ConvNeXtFoldGB.lean.

The gate (lake build convnext-fwd-b-tie): the per-example chain and this one must emit byte-identical forwards, and train steps that differ on the conv-VJP transpose/reverse pair and nothing else (commuting ops on disjoint axes). That is available because every batched form was built to emit its per-example peer's text byte-for-byte (tests/TestBatchedEmitTie.lean checks each form). So the whole-net statement is the per-form statement composed — and if it ever fails, the tie file localises which form did it in one run. The committed bytes are this chain's, so the gate renders the per-example chain and compares it against them.

The batched shape helpers #

reassoc/unassoc in the per-example renderer cast SHlo (c*h*h) ↔ SHlo (c*(h*h)). At the batched index the same cast has to happen UNDER N * ·, which is bnBatchLA's existing move (congrArg (N * ·) (Nat.mul_assoc …)). It is a reindex, not a function change — the emitted text is unaffected, since skel never sees the index.

def Proofs.StableHLO.reassocB {N c h : ℕ} (e : SHlo (N * (c * h * h))) :
SHlo (N * (c * (h * h)))
Equations
Instances For

    ConvNeXt-T / ImageNet's per-replica batch, spelled ONCE for the whole convnextin_* family.

    64, not the 32 every other ConvNeXt render uses, and the reason is a PAIRING fact rather than a throughput one. This net's JAX reference trains at global 256 = 4 × 64 (its own banner: batch 256 (4x64) · SPE 5004). At 4 × 32 the verified job would run half the batch, twice the updates, and an LR off the linear-scaling rule: a RECIPE difference, not a lowering difference, and it would forfeit the only thing this pair is in the book for: ConvNeXt has no BatchNorm, so a tie here says the LOWERER is clean and an offset here implicates the lowerer or the feed fleet-wide. Neither reading survives if the batch does not match.

    Spelled once because it has to reach every convnextin_* #eval AND both halves of every train step (body at N := bB, wrapper via cBS).

    ConvNeXt-S's and -B's EMA pair renders and eval forwards are at this batch too; their adam* train steps stay at the 32 default. A job's LEAN_MLIR_BATCH must match its render — the PRECHECK in scripts/jobs/cnx-default-4gpu.conf greps the artifact for its own baked batch and is the template for that.

    Equations
    Instances For

      Eps and the stage table — read from the per-example renderer's own constants where they are public, restated where they are private. Restated, not re-derived: if these drift the byte tie fails loudly, which is the point of tying against the committed artifact rather than against a second copy of the shapes. The batch is a parameter, not one of these constants — see the note above.

      Equations
      Instances For
        Equations
        Instances For

          The sites, one per per-example site in ConvNeXtRender.lean #

          def Proofs.StableHLO.lnFwdSiteB (bB : ℕ) (gN btN xin : String) (c h : ℕ) :

          One channel-LN forward site, batched: transpose to [h·w, c], normalise each spatial row over its channels at the scalar identities %one/%zero, apply the real [c] affine, transpose back. Five batchOps where the per-example peer has five bare nodes.

          m := h*h is the SPATIAL row count PER EXAMPLE and N := bB is the batch. Collapsing those two into one index is exactly the defect this file exists to remove, and it is why the descriptors carry (m, n) internally instead of reading the SHlo index.

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

            The head LN's input-VJP, batched peer.

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

              The head LN's γ / β tails, batched peers — both contract the batch.

              Equations
              Instances For
                def Proofs.StableHLO.fwdBlockB (gf : GeluForm) (bB : ℕ) (pfx xin : String) (c e h : ℕ) (drop : Option ℕ := none) (bf16 : Bool := false) :

                One ConvNeXt block forward, batched: depthwise 7×7 → channel-LN → 1×1 expand → GELU → 1×1 project → LayerScale → [drop] → + skip. The residual add is addVB, the binary batched form.

                drop = some i puts the stochastic-depth site at ramp index i on the residual branch, i.e. between layerScaleCh and the addVB. That placement is the whole correctness question and the obvious gate is blind to it — at an all-ones mask 1 ⊙ (branch + x) = branch + x exactly, so a site on the block OUTPUT is the same function bit-for-bit and every endpoint gate passes on it (measured on EfficientNet). scripts/probes/misplace_drop_sites.py builds exactly that render — same SSA names, order, types and line count — and it is the control that licenses believing any green run here.

                At drop = none no pretty call happens, so the fresh-name counter does not move and the drop-free chain re-renders byte-identically. That is what keeps convnext-fwd-b-tie and the committed artifacts free of this feature.

                Equations
                • One or more equations did not get rendered due to their size.
                Instances For
                  def Proofs.StableHLO.fwdDownB (bB : ℕ) (pfx xin : String) (ci co h2 : ℕ) (bf16 : Bool := false) :

                  One downsample forward, batched: channel-LN then 2×2/s2 conv.

                  Equations
                  • One or more equations did not get rendered due to their size.
                  Instances For
                    def Proofs.StableHLO.convNextFwdChainB (gf : GeluForm) (nClasses : ℕ := 10) (sd : Bool := false) (V : CnxDims := bTiny) (bf16 : Bool := false) (bB : ℕ := 32) (f : ℕ := 7) :

                    The full ConvNeXt-T [3,3,9,3] forward at the batched index. Node for node the same chain convNextFwdChain emits — 4×4/s4 patchify stem (3→96, 224→56) → stem channel-LN → 4 stages at 56/28/14/7 with 2×2/s2 downsamples between → GAP(7×7) → dense(768→nClasses).

                    Every node is a batchOp/*B form, so den is a batchMap/batchMapAux at N := bB and the batch is an index of the AST rather than a number only pretty knows. That is the entire content of the move; the emitted text is unchanged, which the tie checks.

                    sd := true adds the 18 stochastic-depth sites, one per block, at ramp index cnxBlockIdx si j — and the sites are emitted in the FORWARD as well as the train step deliberately: at eval the driver supplies an all-ones mask, so they are the exact identity, and the forward ⊂ train-step prefix audit keeps a partner for the SD render instead of quietly not covering it.

                    Equations
                    • One or more equations did not get rendered due to their size.
                    Instances For
                      def Proofs.StableHLO.convNextFwdRenderB (gf : GeluForm) (funcName : String := "convnext_fwd_b") (nClasses : ℕ := 10) (banner : String := " // ── ConvNeXt-T forward at the BATCHED index N := B: every op is pretty(batchOp …) except the %one/%zero LayerNorm constants ──\n") (sd : Bool := false) (V : CnxDims := bTiny) (bf16 : Bool := false) (bB : ℕ := 32) (s : ℕ := 224) :

                      @convnext_fwd_b — the batched-index peer of convNextFwdText, same signature (182 parameters at ConvNeXt-T) and same %x. This WRITES convnext_fwd, convnextin_fwd, convnextsin_fwd and convnextbin_fwd (the #evals at the bottom of this file); convnext-fwd-b-tie renders the per-example chain against these bytes.

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

                        The backward #

                        Every site below is the per-example site with its node swapped for the batched form. Two of them are where increments 2 and 3's constructors earn their keep, and both are cases the emit tie alone cannot judge:

                        def Proofs.StableHLO.lnBackSiteB (bB : ℕ) (gN xName cot : String) (c h : ℕ) :

                        One channel-LN input-VJP site, batched.

                        Equations
                        • One or more equations did not get rendered due to their size.
                        Instances For
                          def Proofs.StableHLO.lnGammaTailB (bB : ℕ) (_gN xName cot : String) (c h : ℕ) :

                          The γ tail for one LN site — the two-level veclnGammaGradB.

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

                            The β tail — the two-level rowDenseBiasGradB, contracting batch and spatial rows.

                            Equations
                            • One or more equations did not get rendered due to their size.
                            Instances For
                              def Proofs.StableHLO.bwdBlockB (gf : GeluForm) (bB : ℕ) (pfx dy : String) (b : FNames) (c e h : ℕ) (drop : Option ℕ := none) (bf16 : Bool := false) :

                              One ConvNeXt block backward: the cotangent chain only, param tails separate (the per-example renderer factors it the same way, which is what makes ConvNeXt the cheapest of the five to thread).

                              THE DROP'S BACKWARD IS THE SAME OP AT THE SAME MASK (Proofs.dropPath_vjp_is_self): a diagonal linear map is its own transpose, so there is no *Grad peer to keep in step and no second emitter to drift.

                              IT APPLIES TO THE BRANCH ONLY, and that mirrors the forward's placement exactly. The dropped cotangent cotD feeds the whole branch — LayerScale, project, GELU, expand, LN, depthwise — including every parameter gradient computed off it. The skip's fan-in at the bottom keeps the RAW dy. Dropping there too would attenuate the identity path, which is s ⊙ (branch + x) arriving by the other door.

                              AND dyd IS RETURNED BECAUSE ONE PARAMETER GRADIENT READS IT DIRECTLY. LayerScale's γ gradient is Σ (cot ⊙ p) at the cotangent of the LayerScale OUTPUT — which is s ⊙ dy once a drop site sits between LayerScale and the add, not dy. Every other block gradient descends from cot_p and inherits the scale for free; %…lg is the one that would silently be computed against an undropped cotangent. It type-checks, trains and descends: 18 of 180 gradients wrong by a per-example factor, on the parameter stochastic depth is about. At drop = none this is dy itself, so nothing moves.

                              Equations
                              • One or more equations did not get rendered due to their size.
                              Instances For
                                def Proofs.StableHLO.bwdDownB (bB : ℕ) (pfx dy xin : String) (ci co h2 : ℕ) (bf16 : Bool := false) :

                                One downsample backward.

                                Equations
                                • One or more equations did not get rendered due to their size.
                                Instances For
                                  def Proofs.StableHLO.blockParamGradB (bB : ℕ) (pfx : String) (b : FNames) (cot_p cot_e cot_n cot_d dy : String) (c e h : ℕ) (bf16 : Bool := false) :

                                  The parameter gradients of one block — every one a *GradB, i.e. Σ_n over the batch of the per-example gradient on batchSlice n. AdamW only: the SGD-inline tail stays in the per-example renderer, since %lr is a runtime operand on the AdamW path and a baked literal on the SGD one — and ConvNeXtStepTie.lean is stated at those bytes.

                                  Equations
                                  • One or more equations did not get rendered due to their size.
                                  Instances For
                                    def Proofs.StableHLO.downParamGradB (bB : ℕ) (pfx downLn downIn cot_n dy : String) (ci co h2 : ℕ) (bf16 : Bool := false) :

                                    The parameter gradients of one downsample.

                                    Equations
                                    • One or more equations did not get rendered due to their size.
                                    Instances For
                                      def Proofs.StableHLO.convNextBackAllB (gf : GeluForm) (smooth : Option (String × String × String) := none) (nClasses : ℕ := 10) (sd : Bool := false) (V : CnxDims := bTiny) (bf16 : Bool := false) (bB : ℕ := 32) :

                                      The whole-net batched traversal — forward + cotangent + every parameter gradient, the batched peer of convNextBackAll true (some …). Returns (code, gradMap, softmaxSSA) with the same shape, so the AdamW tail in ConvNeXtRender.lean can consume either.

                                      AdamW only, deliberately. The per-example traversal serves both renders off one adam flag; this one does not, because the SGD path bakes lr as a literal where AdamW takes it as a runtime %lr operand, and an SGD render at the batched index is not something any config asks for yet. Adding the flag later is cheap; an artifact nobody loads is a silent-hyperparameter hazard.

                                      The %dgi/%dgb/%dgn/%dgd/%dgapf GAP-backward block is hand-written text on both sides, carried over verbatim. It is one of the declared non-AST carve-outs, so the batched move neither improves nor degrades it — but note it is parameterised by bB and therefore already batch-correct, which is why it needs no peer.

                                      Equations
                                      • One or more equations did not get rendered due to their size.
                                      Instances For
                                        def Proofs.StableHLO.convNextAdamTrainStepBText (gf : GeluForm) (alphaStr negAlphaKStr bStr : String) (replicas : ℕ := 1) (nClasses : ℕ := 10) (slug : String := "convnext") (ema wdExclude : Bool := false) (wdStr : String := "0.0001") (clip : Bool := false) (clipStr : String := "1.0") (sd : Bool := false) (V : CnxDims := bTiny) (bf16 : Bool := false) (bB : ℕ := 32) (acc : ℕ := 1) :

                                        The ConvNeXt-T AdamW train step at the batched index. It is the SAME renderer the per-example path uses — convNextAdamTrainStepText with traversal pointed at convNextBackAllB — not a copy.

                                        That is possible because the AdamW tail is entirely parameter-space: adamMNextF, adamVNextF, adamWParamF, gradSumSqAccF and clipScaleF are indexed by the parameter's own size and never see the batch. So "the AdamW tail at the batched index" is no work at all — the batch is factored out of it by the ops' own shapes. The only thing that moves is which traversal produced the gradients.

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

                                          The drop-free forward's banner — the line ConvNeXtRender.convNextFwdText emits, restated here so the byte tie can demand byte-identity rather than "identical apart from a comment". Worth the parameter: a tie that compares modulo one line is a tie with a hole in it, and the hole is exactly where a renderer's own description of what it did would live. This is the banner the committed convnext_fwd / convnextin_fwd / convnextsin_fwd / convnextbin_fwd carry, which is why it takes the size: the model name is derived from the stage table, exactly as the per-example line derives it.

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

                                            The SD forward's banner. Its own, and not cnxFwdBanner, because these bytes ARE a different render and a banner claiming otherwise would misdescribe the artifact it heads.

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