Documentation

LeanMlir.Proofs.Codegen.ResNet50RenderB

ResNet-50 train step rendered from the verified AST, at the BATCHED index #

The bottleneck peer of ResNet34RenderB, block for block: batch BN (bnBatchF, reduced over [0,2,3]), the whole graph at N := B, the un-fused *GradB parameter gradients, and the proven AdamW / heavy-ball tail — optAllParams/optConstsB come from RenderKit, imported, not copied, so the optimizer has one definition across the nets.

The three block forms, and why the third exists #

ResNet50FullBVJP.lean carries the certified VJPs; this file renders them.

rendererblockwhere
bnkIdFwdB / bnkIdBackGradBidentity, oc → mid → mid → oc12 blocks
bnkProjFwdB / bnkProjBackGradBstride-1 projectionstage 1 block 0 only
bnkStridedFwdB / bnkStridedBackGradBstrided projectionstages 2/3/4 block 0

The stride is on the 3×3 (W2), not the leading 1×1 — v1.5 / torchvision, which is what jax/MainResnet50Imagenet.lean trains. So in the strided block conv1, bn1 and relu1 all run at the input resolution 2hh, and only conv2 decimates. That asymmetry is why the strided renderer carries two sets of zero-vectors (zIn/zMidIn at 2hh, zMid/zOut at hh) where the other two carry one.

mid = oc / 4, so an identity block's mid is a QUARTER of its channel count, not a multiple. VerifiedSpec.bottleneckStageSpec already computes it that way and already selects the projecting form when stride != 1 || ic != oc — which is exactly how stage 1 (64→256 at stride 1) gets a projection. This file's block sequence must agree with that spec's ORDER, because the driver packs [θ|m|v] off net.specs; r50SigList below is the render's side of that contract and #guards pin the counts.

No conv biases #

Every conv here is bias-free (convBnNB in the spec — 9 tensors per identity block, 12 per projection block), matching R34's ImageNet render. The convBias plumbing R34's renderer carries for its CIFAR-era artifacts is deliberately absent rather than threaded and passed false.

Proofs tier, and no hand-written render to tie against #

The Proofs-tier files for this net are ResNet50FullB (the batched ℝ forward and its typed graph), ResNet50StepTieB (every parameter-gradient node pinned to the cotangent the emitted backward chain delivers) and ResNet50SyncStepTieB (the data-parallel sync-BN twin). FwdGraphTextTies checks by #guard that this file's stem, block and head emitters print the same bytes as pretty of the matching typed graphs (r50StemGraphB, r50IdGraphB, r50ProjGraphB, r50DownGraphB, r34HeadGraphB), at batch 2 and one shape each.

Every other net's verified render also has a bit-exact numeric tie against a hand-written artifact. R50 has no such artifact, so no such tie exists for it. The substitutes are the layer-level VJP oracle (tests/vjp_oracle/run.sh) and a keep-1 known-answer check.

The 161 R50 parameter inputs, in func-arg order: stem (3), the [3,4,6,3] bottleneck stages, then the head. Single source for %p, %pm, %pv and the return order.

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 106 running-stat inputs — 53 BN layers × (μ, var), μ and var interleaved per layer, which is how the driver packs runningBnStats off bnChannels.

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

        Saved forward SSA names the bottleneck's backward + gradient passes reference. The 3-conv peer of BFwdB: two interior ReLUs (r1, r2) where the basic block has one.

        Instances For

          R50's stochastic-depth site count — one per bottleneck block, [3,4,6,3] = 16.

          It is the count the RAMP's denominator is read against (totalDrop − 1 = 15), so it is a definition rather than a literal 16 sprinkled at three sites.

          Equations
          Instances For

            The %dp0 … %dp15 mask arguments, tensor<Bxf32> each — appended to the train step's signature only under sd. Mirrors cnxDropSig.

            Equations
            Instances For

              some i under sd, none otherwise — the per-block ramp index handed to the six block emitters. The index is the BLOCK index and the two coincide here, because every R50 bottleneck drops; EfficientNet's do not (its reference advances the ramp counter inside a skip guard), which is why efficientnetVerified.dropKeeps is a literal array and this is a range. At sd := false this is none at every site and NOT ONE pretty call happens, so every committed R50 artifact re-renders byte-identically.

              Equations
              Instances For
                def Proofs.StableHLO.bnkIdFwdB (B mid oc hh : ℕ) (epsStr p xName : String) (bf16 : Bool := false) (drop : Option ℕ := none) (replicas : ℕ := 1) (sync : Bool := false) :

                Identity bottleneck forward: 1×1 → BN → relu → 3×3 → BN → relu → 1×1 → BN → (+x) → relu, all at hh, channels oc → mid → mid → oc.

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

                  Stride-1 projection bottleneck forward — R50 stage 1 block 0, and nowhere else. cin → mid → mid → oc with the resolution unchanged, so the projection is a plain 1×1 conv → BN, NOT the strided one bnkStridedFwdB uses.

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

                    Strided projection bottleneck forward — stages 2/3/4 block 0. cin → mid → mid → oc, 2hh → hh.

                    conv1/bn1/relu1 run at the input resolution 2hh; only conv2 (the 3×3) is strided. v1.5.

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

                      Everything the whole-net render needs out of ONE forward traversal of ResNet-50: the emitted code, the logits, and every saved activation the backward reads.

                      This exists so @resnet50_fwd and @resnet50_adam_train_step cannot be different nets: both are rendered from this one traversal, batch BN throughout (reduce [0,2,3], divisor B·H·W), and scripts/regen_verified_mlir.sh check holds the forward as a byte prefix of the train step.

                      The eval forward is deliberately not moved onto this chain. bnEval reads frozen per-channel statistics as graph inputs, so it is the same arithmetic in either vocabulary and has no BN-world to disagree about, and it is the artifact that scores, so its bytes stay as they are.

                      Instances For
                        def Proofs.StableHLO.r50StemFwdB (B q : ℕ) (epsStr : String) (bf16 : Bool := false) (replicas : ℕ := 1) (sync : Bool := false) :

                        Stem forward at ladder base q (input 32q): 7×7/s2 conv → batch BN → relu → He et al.'s 3×3/s2 max-pool on %x, output at 8q. ResNet-34's stem at another resolution (same record). 2*q1 rather than a name, because convStrided (h := q1) demands its operand at exactly Vec (B*(3*(2*q1)*(2*q1))) and any other spelling of the same number is a different TERM.

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

                          Head forward: GAP(q×q) → dense(2048→nClasses) on the last block's output xName. Returns the text and the GAP and logit names.

                          Equations
                          • One or more equations did not get rendered due to their size.
                          Instances For
                            def Proofs.StableHLO.r50FwdChainB (B nClasses : ℕ) (epsStr : String) (q : ℕ := 7) (bf16 sd : Bool := false) (replicas : ℕ := 1) (sync : Bool := false) :

                            The ResNet-50 forward chain at the BATCHED index — one traversal, consumed by both @resnet50_fwd and the train step that differentiates it.

                            Equations
                            • One or more equations did not get rendered due to their size.
                            Instances For
                              def Proofs.StableHLO.resnet50TrainStepText (B nClasses : ℕ) (epsStr : String) (replicas : ℕ := 1) (opt : OptRecipe := OptRecipe.adamw) (slug : String := "resnet50in") (bce : Bool := false) (wdStr : String := "") (q : ℕ := 7) (wdExclude gradClip : Bool := false) (clipNorm : Float := 1.0) (bf16 ema sd forceSync noSync : Bool := false) :

                              ResNet-50 [3,4,6,3] bottleneck train step, batch-BN, rendered at N := B.

                              161 θ / 161 m / 161 v, %lr/%bc1/%bc2, 106 running-stat slots and %onehot in; 161 θ' / 161 m' / 161 v', %loss/%bc1/%bc2 and 106 batch stats out. Parameter ORDER comes from r50SigList, the same single source the signature and the return list use, so the arity/order contract cannot drift within this file — and it is written to agree with VerifiedSpec.bottleneckStageSpec, which is what the DRIVER walks.

                              The block sequence is [3,4,6,3] with block 0 of every stage projecting. Stage 1's projects at stride 1 (bnkProjFwdB); stages 2/3/4 project strided.

                              Instances For

                                These live here rather than in a ResNet50Render.lean peer of R34's split. That split exists because R34 has a per-example-BN net AND a batch-BN net, and they are different functions, so they get different files. R50 has only the batch-BN net, so one file keeps r50SigList and every consumer of it together.

                                The chain below is the EVAL forward only: every BN site is RenderKit.bnEvalSite (frozen running statistics, the per-example index), shared with ResNet-34's eval chain. The training forward is the batched r50FwdChainB.

                                def Proofs.StableHLO.resnet50FwdText (B nClasses : ℕ) (epsStr : String) (slug : String := "resnet50in") (q : ℕ := 7) (vSuffix : String := "") :

                                @resnet50in_fwd — 162 inputs (%x + 161 params), logits [B, nClasses]. Batch-statistic BN: rendered from r50FwdChainB, the traversal the train step differentiates, so check_adam_prefix holds it as a byte prefix of the train step.

                                r50FwdChain's .train branch is UNUSED by R50 and must stay that way; it survives only because .eval shares the function. Rendering a forward from it reopens the split.

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

                                  @resnet50in_fwd_eval — the inference forward, every BN site reading frozen running stats. 161 params + 106 stat inputs + %x = 268 inputs. This is what the driver scores through, so its BN order must match r50StatSigList, which it does by sharing the chain.

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