Documentation

LeanMlir.Proofs.Codegen.EfficientNetRender.Basic

EfficientNet-B0 train step rendered ENTIRELY from the verified AST (batched) #

The Chapter-7 peer of MobileNetV2RenderB, for the committed full-16-MBConv EfficientNet-B0 (262 params, the real [t,c,n,s,k] B0 spec). EfficientNet emits true batch-norm, which couples the batch — so the whole net lives at the batched index N·(c·h·w) (StableHLO.batchOp/bnBatchF/the batched backward + param-SGD ops, all Item B).

The SE wrinkle (vs MobileNetV2's relu6 blocks). Each MBConv has a squeeze-excite gate x ⊙ sigmoid(dense W₂ (swish (dense W₁ (GAP x)))), and the committed trainer trains all 4 SE dense params. The fused batchOp seBlock / seBackBatched give the forward value + the SE input cotangent but NOT the SE param grads, so the renderer un-fuses SE: it keeps the fused seBlock for the forward out and ADDITIONALLY emits the un-fused gate subnet s = batchOp gap → e1 = batchOp dense W₁ → z = batchOp swish → e2 = batchOp dense W₂ (only to expose s/e1/z/e2); the SE param grads chain seReduceB → sigmoidBack(e2) → denseWeightSgdB/denseBiasSgdB (W₂) → denseRowBack(W₂) → swishBack(e1) → denseWeightSgdB/denseBiasSgdB (W₁), and dx reuses the fused seBackBatched. Activations are swish (smooth, no relu6 kink), the head GAP-back uses the batched gapBackBatched.

Render is value-independent (skel erases values), so placeholder zeros + lr := 0/ε := 0 are passed; the emitted lrStr/epsStr literals carry the real values.

The batched index: why every N below is B, not 1 #

Building the graph at the batch-unit index N = 1 and letting pretty B supply the real batch would be sound for the ops where the batch is a parallel index — batchOp's den is batchMap N (denOp op), which at N = 1 is the per-example op, exactly what the emit applies across the batch. It is not sound for the ops where the batch is a reduction axis: bnBatchF and bnBatchBack reduce μ/var over [0,2,3], and the whole *SgdB param family (bnGammaSgdB, bnBetaSgdB, dense{Weight,Bias}SgdB, conv{,Strided}WeightSgdB, depthwise{,Strided}WeightSgdB) sums the per-example gradient over Fin N. At N = 1 each of those dens would describe a ONE-EXAMPLE function while the emitted text reduces over all B.

So the whole graph sits at N := B, where those dens are honest. The obstacle there is not the batch-coupled ops (their emitters discard N and use B, so they render identically at any N) but the pointwise ones: swishF/swishBack/sigmoidBack/addV/sub carry only the SHlo index and emit tensor<B×n> from it, so at the batched index N·s they emit tensor<B×(N·s)> — which does not even typecheck against its own operand. Hence the batched forms, which all separate the batch N from the per-example emit width n:

softmaxRow's m and denseRowBack's rows are NOT the batch — they are rows per example (ViT uses m := 197 tokens; a classifier head has one logit row), and they stay 1 here.

The emitted text does not depend on N, and it cannot witness the den side — the render is value-independent, so a descriptor holding the wrong saved activation would render exactly the same bytes. That half is carried by the rfl faithfulness theorems in StableHLO.Basic.

Saved forward SSA names a block's backward + SGD passes reference.

  • code : String
  • code : String
  • o : String
  • o : String
  • ec : String
  • ec : String
  • en : String
  • en : String
  • er : String
  • er : String
  • dc : String
  • dc : String
  • dn : String
  • dn : String
  • dr : String
  • dr : String
  • se : String
  • se : String
  • s : String
  • s : String
  • e1 : String
  • e1 : String
  • z : String
  • z : String
  • e2 : String
  • e2 : String
  • pc : String
  • pc : String
  • The block's BN layers in forward order: (BN-input SSA, channels, spatial side). The AdamW render turns each into a bnBatchMeanB/bnBatchVarB pair — the batch statistics a batch-BN train step has to hand back so the host can EMA them into the eval forward's frozen stats. The SGD render has no such outputs and ignores this field.

    Order is expand-BN → depthwise-BN → project-BN, with the expand entry ABSENT for the no-expand block (b1, t = 1), because that is the layout the driver's bnChannels metadata and @efficientnet_fwd_eval read positionally. Getting it wrong is silent: the arities still match and the wrong layer's statistics simply flow into the wrong eval slot.

    The second component is the layer's stat prefix — @efficientnet_fwd_eval takes %{prefix}mu/%{prefix}var there, and the AdamW train step hands the matching batch μ/var back from the SAME entry. So the eval signature, the eval BN sites and the train step's stat outputs all come off this one list; there is deliberately no parallel 49-entry table.

  • The block's BN layers in forward order: (BN-input SSA, channels, spatial side). The AdamW render turns each into a bnBatchMeanB/bnBatchVarB pair — the batch statistics a batch-BN train step has to hand back so the host can EMA them into the eval forward's frozen stats. The SGD render has no such outputs and ignores this field.

    Order is expand-BN → depthwise-BN → project-BN, with the expand entry ABSENT for the no-expand block (b1, t = 1), because that is the layout the driver's bnChannels metadata and @efficientnet_fwd_eval read positionally. Getting it wrong is silent: the arities still match and the wrong layer's statistics simply flow into the wrong eval slot.

    The second component is the layer's stat prefix — @efficientnet_fwd_eval takes %{prefix}mu/%{prefix}var there, and the AdamW train step hands the matching batch μ/var back from the SAME entry. So the eval signature, the eval BN sites and the train step's stat outputs all come off this one list; there is deliberately no parallel 49-entry table.

  • stE : String
  • stE : String
  • stD : String
  • stD : String
  • stP : String
  • stP : String
  • bnSt : List String
  • bnSt : List String
Instances For
    Instances For

      Every leaf of the backward ends at a parameter, and there are exactly two things it can emit: the un-fused gradient (adam := true) or the fused SGD update θ − lr·g (false). The *SgdB_eq_grad theorems say den (xSgdB …) = θ − lr · den (xGradB …) by rfl, and tests/TestBatchedEmitTie.lean checks the emit side of the same statement: each *GradB render is a byte-PREFIX of its *SgdB peer's, the tail being exactly the const-lr / multiply / subtract.

      These six helpers are what let ONE backward traversal serve both renders. The alternative — a second copy of the 16-MBConv backward for AdamW — is the double-writer disease one level down, in code rather than in artifacts: two emitters for one step that can compute different functions.

      lrStr is threaded but unused in adam mode: the AdamW render's learning rate is the runtime %lr argument, not a baked literal. The placeholder values are irrelevant either way — skel erases them, so the render is value-independent.

      SE forward on c channels at hh×ww, reduce dim r: one fused seBlock, whose text computes the squeeze (GAP), reduce dense (W₁), swish and excite dense (W₂) on the way to the gate. Their SSA names (seBlockSavedNames) are the activations the SE backward reads, so the forward computes the SE once — it is pretty of the typed graph's seBlock node — rather than a second, un-fused copy beside it. Returns (code, s, e1, z, e2, seOut).

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

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

        Equations
        Instances For

          The block indices (0-based) that carry a drop site, in signature order. This is the single source: enetDropSig maps over it to build the %dp<i> inputs, the traversal passes the same i at each eFwd call site, and the driver reads it to know how many scales to supply and at which ramp index. The two routes fail LOUDLY if they disagree — an entry with no call site leaves an unused input (arity mismatch at the driver), and a call site with no entry emits an undeclared %dp<i> (rejected by the lowerer). Neither is silent.

          Equations
          Instances For

            The number of per-example drop-path scale inputs a stochastic-depth render takes.

            Equations
            Instances For

              The %dp<i>: tensor<Bxf32> inputs, appended to a render's signature when stochastic depth is on. Empty when off.

              Equations
              Instances For

                The classifier's input width — EfficientNet-B0's head channel count, i.e. the GAP output and hence the dropout mask's per-example width. Independent of nClasses.

                Equations
                Instances For

                  The %do: tensor<B×1280xf32> input, appended when classifier dropout is on. Empty when off, which is what keeps the inertness gate byte-identical.

                  It goes after enetDropSig, i.e. dead last in every signature. Two independent reasons, and the second is the one that bites: a parameter inserted mid-list captures an existing positional slot and the driver walks these signatures positionally; and the drop-mask tail is what the DP shim shards by COUNT from the end (n_shard_tail), so a per-example input placed before them would be counted as one of them.

                  Equations
                  • One or more equations did not get rendered due to their size.
                  Instances For
                    def Proofs.StableHLO.eFwdBody (B ic mid oc hh kd r : ℕ) (mode : BnMode) (epsStr p xName : String) (convBias : Bool) (bf16 : Bool := false) (replicas : ℕ := 1) (sync : Bool := false) :

                    Stride-1 expand MBConv forward body (shared by residual + no-skip): expand 1×1 conv-bn-swish → depthwise(kd) conv-bn-swish → SE → project 1×1 conv-bn. Returns the EFwd WITHOUT the final residual (caller adds the addV for residual blocks).

                    Equations
                    • One or more equations did not get rendered due to their size.
                    Instances For
                      def Proofs.StableHLO.eFwd (B ic mid oc hh kd r : ℕ) (mode : BnMode) (epsStr p xName : String) (convBias : Bool) (drop : Option ℕ := none) (bf16 : Bool := false) (replicas : ℕ := 1) (sync : Bool := false) :

                      Residual stride-1 MBConv forward (ic = oc): body + addV skip.

                      Equations
                      • One or more equations did not get rendered due to their size.
                      Instances For
                        def Proofs.StableHLO.eFwdNoSkip (B ic mid oc hh kd r : ℕ) (mode : BnMode) (epsStr p xName : String) (convBias : Bool) (bf16 : Bool := false) (replicas : ℕ := 1) (sync : Bool := false) :

                        No-skip stride-1 MBConv forward (ic ≠ oc, b9/b16): body, output = project-BN out.

                        Equations
                        Instances For
                          def Proofs.StableHLO.eFwdStrided (B ic mid oc hh kd r : ℕ) (mode : BnMode) (epsStr p xName : String) (convBias : Bool) (bf16 : Bool := false) (replicas : ℕ := 1) (sync : Bool := false) :

                          Strided MBConv forward (b2/b4/b6/b12): expand at the input 2hh×2ww, depthwise downsamples 2hh×2ww → hh×ww, project 1×1 at hh×ww. NO skip.

                          Equations
                          • One or more equations did not get rendered due to their size.
                          Instances For
                            def Proofs.StableHLO.eFwdNoExp (B ic oc hh kd r : ℕ) (mode : BnMode) (epsStr p xName : String) (convBias : Bool) (bf16 : Bool := false) (replicas : ℕ := 1) (sync : Bool := false) :

                            No-expand MBConv forward (b1, t=1): depthwise(kd, on ic channels)-bn-swish → SE → project 1×1 (ic→oc)-bn. NO expand, NO skip. ec/en unused; er = block input (= depthwise input).

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

                              Names are stored WITHOUT the leading % and shapes as List Nat rather than a rendered tensor<…> string, because the AdamW render needs both forms: %{nm}/%{nm}m/%{nm}v for the moment slots, and the raw dimensions for the emitted Adam ops (adamMNextF's ds); ty ds renders the type string.

                              def Proofs.StableHLO.enetSig (nClasses : ℕ) (convBias : Bool) :

                              Full 262-param EfficientNet-B0 signature, func-arg order: stem(4) + b1 no-exp(12) + b2..b16 expand(15×16) + head(4) + dense(2) = 4+12+240+4+2 = 262 tensors.

                              Instances For

                                Every channel width EfficientNet-B0 uses as a conv bias — stem 32, each block's mid and oc off the [t,c,n,s,k] table, head 1280. One list feeding all four zeroBiasPrelude calls, so a convBias := false render cannot declare the constants in one artifact and not another.

                                NOT the SE widths. SE's two 1×1 convs are followed by an ACTIVATION, not BN, so nothing absorbs their biases and the reference carries them; they stay real parameters and are never bound to a zero constant. The rule for which biases fold — a rank-1 parameter after a rank-4 kernel — excludes them because SE's params are rank-2.

                                Equations
                                Instances For

                                  Every SSA name the EfficientNet-B0 forward produces, plus the 49-entry BN stat layout. efficientnetFwd{,Eval}FaithfulV return just logits; the train steps additionally consume the stem/head names and the 16 block records on the way back.

                                  • code : String
                                  • code : String
                                  • stc : String
                                  • stc : String
                                  • stn : String
                                  • stn : String
                                  • str : String
                                  • str : String
                                  • blocks : Array EFwd
                                  • blocks : Array EFwd
                                  • hc : String
                                  • hc : String
                                  • hn : String
                                  • hn : String
                                  • hr : String
                                  • hr : String
                                  • gap : String
                                  • gap : String
                                  • cin : String

                                    THE CLASSIFIER'S ACTUAL INPUT — gap with classifier dropout OFF, the dropoutB output with it ON. It exists as its own field, rather than every consumer reading gap, because there are TWO consumers and one of them is easy to miss:

                                    • the dense forward, which obviously reads it; and
                                    • the dense WEIGHT gradient, ∂L/∂W = Σ_b dy_b ⊗ (input_b) — which reads the dense's input, i.e. the DROPPED activation, not the pooled one.

                                    Feeding dnW the undropped gap type-checks, trains, descends, and is wrong on the one parameter dropout acts through. It is invisible to every ones-mask gate this feature has, because at mask ≡ 1 the two values are equal. It is a named field so that every consumer of the value dropout displaces reads it by name.

                                  • cin : String

                                    THE CLASSIFIER'S ACTUAL INPUT — gap with classifier dropout OFF, the dropoutB output with it ON. It exists as its own field, rather than every consumer reading gap, because there are TWO consumers and one of them is easy to miss:

                                    • the dense forward, which obviously reads it; and
                                    • the dense WEIGHT gradient, ∂L/∂W = Σ_b dy_b ⊗ (input_b) — which reads the dense's input, i.e. the DROPPED activation, not the pooled one.

                                    Feeding dnW the undropped gap type-checks, trains, descends, and is wrong on the one parameter dropout acts through. It is invisible to every ones-mask gate this feature has, because at mask ≡ 1 the two values are equal. It is a named field so that every consumer of the value dropout displaces reads it by name.

                                  • logits : String
                                  • logits : String
                                  • The 49 BN layers as (BN-input SSA, stat prefix, channels, spatial side), stem → blocks in forward order → head. Single source for the eval signature, the eval BN sites and the AdamW train step's returned batch statistics — see EFwd.bns.

                                  • The 49 BN layers as (BN-input SSA, stat prefix, channels, spatial side), stem → blocks in forward order → head. Single source for the eval signature, the eval BN sites and the AdamW train step's returned batch statistics — see EFwd.bns.

                                  • sst : String
                                  • sst : String
                                  • hst : String
                                  • hst : String
                                  • bnSt : List String
                                  • bnSt : List String
                                  Instances For

                                    The stem's saved SSA names: conv, BN, BN stats ("" at one replica), swish output.

                                    Instances For
                                      def Proofs.StableHLO.enetStemFwdB (B : ℕ) (mode : BnMode) (epsStr : String) (convBias : Bool) (bf16 : Bool := false) (replicas : ℕ := 1) (sync : Bool := false) :

                                      Stem forward: 3×3/s2 XLA-SAME conv (3→32, 224→112) → BN → swish, on %x.

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

                                        The head's saved SSA names up to the GAP: conv, BN, BN stats, swish, GAP.

                                        Instances For
                                          def Proofs.StableHLO.enetHeadFwdB (B _nClasses : ℕ) (mode : BnMode) (epsStr xName : String) (convBias : Bool) (bf16 : Bool := false) (replicas : ℕ := 1) (sync : Bool := false) :

                                          Head forward up to the GAP: 1×1 conv (320→1280) → BN → swish → GAP(7×7), on block 16's output xName. The classifier dropout and the dense stay in the chain (cd decides the dropout).

                                          Equations
                                          • One or more equations did not get rendered due to their size.
                                          Instances For
                                            def Proofs.StableHLO.efficientnetFwdFaithfulV (B nClasses : ℕ) (epsStr : String) (convBias : Bool := false) (slug : String := "efficientnet") (sd cd : Bool := false) :

                                            @efficientnet_fwd rendered ENTIRELY from the verified AST — 263 inputs (%x plus the 262 params in enetSig order), returning logits [B, nClasses]. Shares enetFwdChain with the train step, so it is a byte-identical PREFIX of efficientnet_train_step.mlir, ending exactly where the loss begins. Replaces the hand-written emitter in tests/TestEfficientNetFwd.lean.

                                            Equations
                                            • One or more equations did not get rendered due to their size.
                                            Instances For
                                              def Proofs.StableHLO.efficientnetFwdEvalFaithfulV (B nClasses : ℕ) (epsStr : String) (convBias : Bool := false) (slug : String := "efficientnet") (sd cd : Bool := false) :

                                              @efficientnet_fwd_eval rendered ENTIRELY from the verified AST — the inference forward, every BN site consuming frozen per-channel running stats (the bnEval descriptor, den = batchMap N bnPerChannelEvalTensor3) instead of reducing statistics out of its activation. Same 262 params in the same order, plus the 98 stat inputs (49 BN layers × μ/var, interleaved per layer in bnChannels order): 361 inputs.

                                              This is the eval partner of efficientnet_adam_train_step, whose returned batch μ/var the driver EMAs into exactly these slots — and both sides of that contract now come off one bns list rather than two independently-written ones.

                                              Equations
                                              • One or more equations did not get rendered due to their size.
                                              Instances For
                                                def Proofs.StableHLO.efficientnetTrainStepFaithfulV (B nClasses : ℕ) (epsStr lrStr : String) (funcName : String := "efficientnet_train_step") (convBias : Bool := false) :

                                                EfficientNet-B0 (full 16-MBConv) SGD train step rendered ENTIRELY from the verified AST, at the batched index N·(c·h·w). Every emitted line is pretty of a verified SHlo node. Strided stem 3×3/s2 (3→32, 224→112) → b1 (no-expand) → b2..b16 (4 strided downsamples 112→7, 9 residual skips, 2 no-skip widenings) → 1×1 conv-bn-swish head (320→1280) → GAP → dense (1280→nClasses).

                                                The cotangent is plain softmax − onehot with the batch mean folded into lrStr — so the committed lrStr = 0.05 is an effective 1.6 on the mean loss. That is a tuned value, not a slip; the AdamW render below spells the mean explicitly instead.

                                                Equations
                                                • One or more equations did not get rendered due to their size.
                                                Instances For
                                                  def Proofs.StableHLO.enetAdamVariant (B replicas : ℕ) (opt : OptKind := OptKind.adamw) (ema sd cd bf16 wx : Bool := false) (epsMarker : String := "") :

                                                  The driver's variant slug for a given (B, replicas): the artifact is verified_mlir/efficientnet_<variant>_train_step.mlir, the entry point is @efficientnet_<variant>_train_step, and LEAN_MLIR_VARIANT selects it.

                                                  All three must agree, or the shim refuses the call outright ("entry mismatch") rather than running the wrong graph. Deriving the name here is what stops it drifting from the #eval paths below; the #guards at the bottom pin those literal paths against this function.

                                                  B = 32 is deliberately unsuffixed, so the two existing artifacts keep their names and bytes. Same convention as r34AdamVariant.

                                                  Equations
                                                  • One or more equations did not get rendered due to their size.
                                                  Instances For
                                                    def Proofs.StableHLO.efficientnetAdamTrainStepFaithful (B nClasses : ℕ) (epsStr alphaStr negAlphaKStr bStr : String) (replicas : ℕ := 1) (convBias : Bool := false) (slug : String := "efficientnet") (opt : OptKind := OptKind.adamw) (ema sd cd bf16 forceSync wdExclude : Bool := false) :

                                                    EfficientNet-B0 AdamW train step rendered from the verified AST. The certified peer of the hand-written tests/TestEfficientNetTrain.lean render that efficientnet-verified-adam has been training on.

                                                    Same backward as efficientnet_train_step (enetBackAll, one traversal) but taking the un-fused gradients, each fed to the proven AdamW triple. Two things differ from the SGD render and both are load-bearing:

                                                    • the cotangent is label-smoothed (α = 0.1, K = nClasses) with an explicit ÷B, where the SGD render is plain CE with the mean folded into lr. Measured against the hand-written AdamW emitter, this is the same gap ViT had; get it wrong and the tie fails in a way that looks like a bug in the gradient ops.
                                                    • it returns the BN running statistics — batch μ/var per BN layer, bnBatchMeanB/ bnBatchVarB recomputed from that layer's BN input — which the host EMAs into @efficientnet_fwd_eval's frozen stats. The SGD render has no such outputs.

                                                    Interface: 889 in (%x, 262 θ, 262 m, 262 v, %lr/%bc1/%bc2, 98 running-stat slots, %onehot) / 887 out (262 θ', 262 m', 262 v', %loss/%bc1/%bc2, 98 batch stats) — positionally identical to the hand-written render, so trainAdamSched's packed [θ|m|v] protocol is unchanged.

                                                    Unlike ViT's, this tie can pin the forward bit-exactly: EfficientNet has BatchNorm, so the returned batch statistics are a whole-net forward fingerprint no gradient touches.

                                                    Instances For