Documentation

LeanMlir.Proofs.Nets.EfficientNet.EfficientNetChainClose

EfficientNet Item D — the batched backward (cotangent) math, step by step #

The batched analogue of MobileNetV2ChainClose (and of ResNet-34's per-example close, deleted 2026-09-08). The forward graph (Item A, EfficientNetRenderPC.lean) lives at the batched index N·(c·h·w); here we sort the backward math at that same index — proving the per-block gradient (HasVJP) by composing the proven per-op VJPs, lifted to the batch.

The genuinely-new lemma is batchMap_has_vjp: a batch-separable op batchMap N f (every spatial op in the forward graph) has a block-diagonal VJP — f's VJP applied per example. This is what lets seBlockFull_has_vjp, the conv/depthwise/dense VJPs, etc. lift from one example to the whole batch. Mechanically it reuses the existing row-wise machinery: batchMap N f IS Mat.flatten ∘ (apply f to each row) ∘ Mat.unflatten, so rowwise_has_vjp_mat + hasVJPMat_to_hasVJP (Tensor.lean) close it.

(The one batch-coupled op, true batch-norm, is handled separately by the proven bnBatchTensor4_has_vjp — it is NOT a batchMap. swish/sigmoid are pointwise, so swish_has_vjp/sigmoid_has_vjp apply directly at the batched index. Step-by-step, per the plan: batchMap VJP first, then bnBatchLA, then the per-block chains.)

theorem Proofs.batchMap_eq_rowwiseFlat {N a b : } (f : Vec aVec b) :
StableHLO.batchMap N f = fun (v : Vec (N * a)) => Mat.flatten ((fun (A : Mat N a) (r : Fin N) => f (A r)) (Mat.unflatten v))

batchMap N f is the flattened row-wise application of f. Reading the output at flat index idx (decoding to example m, coord c) gives f (row m of the input) c on both sides — the Mat.flatten/unflatten row-major convention is exactly batchMap's finProdFinEquiv split.

batchMap N f is differentiable when f is — it is f applied independently per example.

noncomputable def Proofs.batchMap_has_vjp {N a b : } (f : Vec aVec b) (hf : HasVJP f) (hf_diff : Differentiable f) :

batchMap N f VJP — block-diagonal (the genuinely-new lemma). A batch-separable op's VJP applies f's proven VJP independently per example. The backward, like the forward, reshapes to [N, ·] and runs f.backward row-wise. Reuses rowwise_has_vjp_mat + hasVJPMat_to_hasVJP. This is seBlockFull_has_vjp / the conv-depthwise-dense VJPs "lifted by batchMap" to the whole batch.

Equations
Instances For
    noncomputable def Proofs.reindex_has_vjp {a b : } (σ : Fin bFin a) :

    Generic reindex VJP. reindexCLM σ (gather y ↦ y ∘ σ) is linear; its backward scatters each output cotangent back to the inputs that map to it (the adjoint). Generalizes the manual reindex VJPs (broadcastFlat_has_vjp, bnchwFwd/Back_has_vjp).

    Equations
    Instances For
      theorem Proofs.bnBatchLA_eq_comp (N oc h w : ) (ε : ) (γ β : Vec oc) :
      StableHLO.bnBatchLA N oc h w ε γ β = (reindexCLM (Fin.cast )) bnBatchTensor4 N oc h w ε γ β (reindexCLM (Fin.cast ))

      bnBatchLA is the proven bnBatchTensor4, conjugated by the mul_assoc reindex. Both reindex maps are reindexCLM (Fin.cast …); the middle is the genuinely batch-coupled true batch-norm.

      theorem Proofs.bnBatchLA_differentiable (N oc h w : ) (ε : ) ( : 0 < ε) (γ β : Vec oc) :
      noncomputable def Proofs.bnBatchLA_has_vjp (N oc h w : ) (ε : ) ( : 0 < ε) (γ β : Vec oc) :
      HasVJP (StableHLO.bnBatchLA N oc h w ε γ β)

      True batch-norm VJP at the network's flat index. bnBatchLA's backward is the proven bnBatchTensor4 VJP (batch-coupled — NOT a batchMap), conjugated by the reindex isos.

      Equations
      • One or more equations did not get rendered due to their size.
      Instances For
        noncomputable def Proofs.flatConv_has_vjp {ic oc h w kH kW : } (W : Kernel4 oc ic kH kW) (b : Vec oc) :

        flatConv W b VJP — the per-example 1×1/3×3 conv input-VJP (the HasVJP3-bridged conv2d).

        Equations
        Instances For
          theorem Proofs.bnSwishStage_differentiable (N : ) {a oc h w : } (op : Vec aVec (oc * h * w)) (hop : Differentiable op) (ε : ) ( : 0 < ε) (γ β : Vec oc) :
          Differentiable (swish (N * (oc * h * w)) StableHLO.bnBatchLA N oc h w ε γ β StableHLO.batchMap N op)

          Differentiability of a batched conv/depthwise → bn → swish stage.

          noncomputable def Proofs.bnSwishStage_has_vjp (N : ) {a oc h w : } (op : Vec aVec (oc * h * w)) (hop : Differentiable op) (hopv : HasVJP op) (ε : ) ( : 0 < ε) (γ β : Vec oc) :
          HasVJP (swish (N * (oc * h * w)) StableHLO.bnBatchLA N oc h w ε γ β StableHLO.batchMap N op)

          VJP of a batched conv/depthwise → bn → swish stage: lift op's VJP per example, then the proven true-BN VJP, then the pointwise swish VJP.

          Equations
          • One or more equations did not get rendered due to their size.
          Instances For
            theorem Proofs.bnStage_differentiable (N : ) {a oc h w : } (op : Vec aVec (oc * h * w)) (hop : Differentiable op) (ε : ) ( : 0 < ε) (γ β : Vec oc) :

            Differentiability of a batched conv → bn stage (project bottleneck, no swish).

            noncomputable def Proofs.bnStage_has_vjp (N : ) {a oc h w : } (op : Vec aVec (oc * h * w)) (hop : Differentiable op) (hopv : HasVJP op) (ε : ) ( : 0 < ε) (γ β : Vec oc) :

            VJP of a batched conv → bn stage.

            Equations
            • One or more equations did not get rendered due to their size.
            Instances For
              theorem Proofs.cbsB_differentiable (N : ) {ic oc h w kH kW : } (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ) ( : 0 < ε) (γ β : Vec oc) :
              Differentiable (cbsB N W b ε γ β)
              noncomputable def Proofs.cbsB_has_vjp (N : ) {ic oc h w kH kW : } (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ) ( : 0 < ε) (γ β : Vec oc) :
              HasVJP (cbsB N W b ε γ β)
              Equations
              Instances For
                theorem Proofs.stemB_differentiable (N : ) {ic oc h w kH kW : } (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ) ( : 0 < ε) (γ β : Vec oc) :
                Differentiable (stemB N W b ε γ β)
                noncomputable def Proofs.stemB_has_vjp (N : ) {ic oc h w kH kW : } (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ) ( : 0 < ε) (γ β : Vec oc) :
                HasVJP (stemB N W b ε γ β)
                Equations
                Instances For
                  theorem Proofs.dwbsB_differentiable (N : ) {c h w kH kW : } (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ) ( : 0 < ε) (γ β : Vec c) :
                  Differentiable (dwbsB N W b ε γ β)
                  noncomputable def Proofs.dwbsB_has_vjp (N : ) {c h w kH kW : } (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ) ( : 0 < ε) (γ β : Vec c) :
                  HasVJP (dwbsB N W b ε γ β)
                  Equations
                  Instances For
                    theorem Proofs.dwbsSB_differentiable (N : ) {c h w kH kW : } (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ) ( : 0 < ε) (γ β : Vec c) :
                    Differentiable (dwbsSB N W b ε γ β)
                    noncomputable def Proofs.dwbsSB_has_vjp (N : ) {c h w kH kW : } (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ) ( : 0 < ε) (γ β : Vec c) :
                    HasVJP (dwbsSB N W b ε γ β)
                    Equations
                    Instances For
                      theorem Proofs.seB_differentiable (N : ) {c h w r : } (W₁ : Mat c r) (b₁ : Vec r) (W₂ : Mat r c) (b₂ : Vec c) :
                      Differentiable (seB N W₁ b₁ W₂ b₂)
                      noncomputable def Proofs.seB_has_vjp (N : ) {c h w r : } (W₁ : Mat c r) (b₁ : Vec r) (W₂ : Mat r c) (b₂ : Vec c) :
                      HasVJP (seB N W₁ b₁ W₂ b₂)
                      Equations
                      Instances For
                        theorem Proofs.projB_differentiable (N : ) {ic oc h w kH kW : } (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ) ( : 0 < ε) (γ β : Vec oc) :
                        Differentiable (projB N W b ε γ β)
                        noncomputable def Proofs.projB_has_vjp (N : ) {ic oc h w kH kW : } (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ) ( : 0 < ε) (γ β : Vec oc) :
                        HasVJP (projB N W b ε γ β)
                        Equations
                        Instances For
                          theorem Proofs.mbNoExpFwdB_differentiable (N : ) {ic oc h w kHd kWd r : } (Wd : DepthwiseKernel ic kHd kWd) (bd : Vec ic) (εd : ) (hεd : 0 < εd) (γd βd : Vec ic) (Wz₁ : Mat ic r) (bz₁ : Vec r) (Wz₂ : Mat r ic) (bz₂ : Vec ic) (Wp : Kernel4 oc ic 1 1) (bp : Vec oc) (εp : ) (hεp : 0 < εp) (γp βp : Vec oc) :
                          Differentiable (mbNoExpFwdB N Wd bd εd γd βd Wz₁ bz₁ Wz₂ bz₂ Wp bp εp γp βp)

                          MBConv1 (no expand) gradient. dw-bn-swish → SE → project-bn.

                          noncomputable def Proofs.mbNoExpFwdB_has_vjp (N : ) {ic oc h w kHd kWd r : } (Wd : DepthwiseKernel ic kHd kWd) (bd : Vec ic) (εd : ) (hεd : 0 < εd) (γd βd : Vec ic) (Wz₁ : Mat ic r) (bz₁ : Vec r) (Wz₂ : Mat r ic) (bz₂ : Vec ic) (Wp : Kernel4 oc ic 1 1) (bp : Vec oc) (εp : ) (hεp : 0 < εp) (γp βp : Vec oc) :
                          HasVJP (mbNoExpFwdB N Wd bd εd γd βd Wz₁ bz₁ Wz₂ bz₂ Wp bp εp γp βp)
                          Equations
                          • One or more equations did not get rendered due to their size.
                          Instances For
                            theorem Proofs.mbStridedFwdB_differentiable (N : ) {ic mid oc h w kHd kWd r : } (We : Kernel4 mid ic 1 1) (be : Vec mid) (εe : ) (hεe : 0 < εe) (γe βe : Vec mid) (Wd : DepthwiseKernel mid kHd kWd) (bd : Vec mid) (εd : ) (hεd : 0 < εd) (γd βd : Vec mid) (Wz₁ : Mat mid r) (bz₁ : Vec r) (Wz₂ : Mat r mid) (bz₂ : Vec mid) (Wp : Kernel4 oc mid 1 1) (bp : Vec oc) (εp : ) (hεp : 0 < εp) (γp βp : Vec oc) :
                            Differentiable (mbStridedFwdB N We be εe γe βe Wd bd εd γd βd Wz₁ bz₁ Wz₂ bz₂ Wp bp εp γp βp)

                            MBConv6 strided gradient. expand-bn-swish → strided dw-bn-swish → SE → project-bn.

                            noncomputable def Proofs.mbStridedFwdB_has_vjp (N : ) {ic mid oc h w kHd kWd r : } (We : Kernel4 mid ic 1 1) (be : Vec mid) (εe : ) (hεe : 0 < εe) (γe βe : Vec mid) (Wd : DepthwiseKernel mid kHd kWd) (bd : Vec mid) (εd : ) (hεd : 0 < εd) (γd βd : Vec mid) (Wz₁ : Mat mid r) (bz₁ : Vec r) (Wz₂ : Mat r mid) (bz₂ : Vec mid) (Wp : Kernel4 oc mid 1 1) (bp : Vec oc) (εp : ) (hεp : 0 < εp) (γp βp : Vec oc) :
                            HasVJP (mbStridedFwdB N We be εe γe βe Wd bd εd γd βd Wz₁ bz₁ Wz₂ bz₂ Wp bp εp γp βp)
                            Equations
                            • One or more equations did not get rendered due to their size.
                            Instances For
                              theorem Proofs.mbResidFwdB_differentiable (N : ) {c mid h w kHd kWd r : } (We : Kernel4 mid c 1 1) (be : Vec mid) (εe : ) (hεe : 0 < εe) (γe βe : Vec mid) (Wd : DepthwiseKernel mid kHd kWd) (bd : Vec mid) (εd : ) (hεd : 0 < εd) (γd βd : Vec mid) (Wz₁ : Mat mid r) (bz₁ : Vec r) (Wz₂ : Mat r mid) (bz₂ : Vec mid) (Wp : Kernel4 c mid 1 1) (bp : Vec c) (εp : ) (hεp : 0 < εp) (γp βp : Vec c) :
                              Differentiable (mbResidFwdB N We be εe γe βe Wd bd εd γd βd Wz₁ bz₁ Wz₂ bz₂ Wp bp εp γp βp)

                              MBConv6 residual gradient. x + (project-bn ∘ SE ∘ dw-bn-swish ∘ expand-bn-swish)(x).

                              noncomputable def Proofs.mbResidFwdB_has_vjp (N : ) {c mid h w kHd kWd r : } (We : Kernel4 mid c 1 1) (be : Vec mid) (εe : ) (hεe : 0 < εe) (γe βe : Vec mid) (Wd : DepthwiseKernel mid kHd kWd) (bd : Vec mid) (εd : ) (hεd : 0 < εd) (γd βd : Vec mid) (Wz₁ : Mat mid r) (bz₁ : Vec r) (Wz₂ : Mat r mid) (bz₂ : Vec mid) (Wp : Kernel4 c mid 1 1) (bp : Vec c) (εp : ) (hεp : 0 < εp) (γp βp : Vec c) :
                              HasVJP (mbResidFwdB N We be εe γe βe Wd bd εd γd βd Wz₁ bz₁ Wz₂ bz₂ Wp bp εp γp βp)
                              Equations
                              • One or more equations did not get rendered due to their size.
                              Instances For
                                theorem Proofs.headFwdB_differentiable (N : ) {c oc h w nC : } (Wh : Kernel4 oc c 1 1) (bh : Vec oc) (εh : ) (hεh : 0 < εh) (γh βh : Vec oc) (Wfc : Mat oc nC) (bfc : Vec nC) :
                                Differentiable (headFwdB N Wh bh εh γh βh Wfc bfc)

                                Head gradient. 1×1 conv-bn-swish → GAP → dense.

                                noncomputable def Proofs.headFwdB_has_vjp (N : ) {c oc h w nC : } (Wh : Kernel4 oc c 1 1) (bh : Vec oc) (εh : ) (hεh : 0 < εh) (γh βh : Vec oc) (Wfc : Mat oc nC) (bfc : Vec nC) :
                                HasVJP (headFwdB N Wh bh εh γh βh Wfc bfc)
                                Equations
                                • One or more equations did not get rendered due to their size.
                                Instances For
                                  noncomputable def Proofs.efficientnetForwardB_has_vjp (N : ) (Ws : Kernel4 32 3 3 3) (bs : Vec 32) (εs : ) (hεs : 0 < εs) (γs βs : Vec 32) (Wd1 : DepthwiseKernel 32 3 3) (bd1 : Vec 32) (εd1 : ) (hεd1 : 0 < εd1) (γd1 βd1 : Vec 32) (Wz1a : Mat 32 8) (bz1a : Vec 8) (Wz1b : Mat 8 32) (bz1b : Vec 32) (Wp1 : Kernel4 16 32 1 1) (bp1 : Vec 16) (εp1 : ) (hεp1 : 0 < εp1) (γp1 βp1 : Vec 16) (We2 : Kernel4 96 16 1 1) (be2 : Vec 96) (εe2 : ) (hεe2 : 0 < εe2) (γe2 βe2 : Vec 96) (Wd2 : DepthwiseKernel 96 3 3) (bd2 : Vec 96) (εd2 : ) (hεd2 : 0 < εd2) (γd2 βd2 : Vec 96) (Wz2a : Mat 96 4) (bz2a : Vec 4) (Wz2b : Mat 4 96) (bz2b : Vec 96) (Wp2 : Kernel4 24 96 1 1) (bp2 : Vec 24) (εp2 : ) (hεp2 : 0 < εp2) (γp2 βp2 : Vec 24) (We3 : Kernel4 144 24 1 1) (be3 : Vec 144) (εe3 : ) (hεe3 : 0 < εe3) (γe3 βe3 : Vec 144) (Wd3 : DepthwiseKernel 144 5 5) (bd3 : Vec 144) (εd3 : ) (hεd3 : 0 < εd3) (γd3 βd3 : Vec 144) (Wz3a : Mat 144 6) (bz3a : Vec 6) (Wz3b : Mat 6 144) (bz3b : Vec 144) (Wp3 : Kernel4 24 144 1 1) (bp3 : Vec 24) (εp3 : ) (hεp3 : 0 < εp3) (γp3 βp3 : Vec 24) (Wh : Kernel4 1280 24 1 1) (bh : Vec 1280) (εh : ) (hεh : 0 < εh) (γh βh : Vec 1280) (Wfc : Mat 1280 10) (bfc : Vec 10) :
                                  HasVJP (headFwdB N Wh bh εh γh βh Wfc bfc mbResidFwdB N We3 be3 εe3 γe3 βe3 Wd3 bd3 εd3 γd3 βd3 Wz3a bz3a Wz3b bz3b Wp3 bp3 εp3 γp3 βp3 mbStridedFwdB N We2 be2 εe2 γe2 βe2 Wd2 bd2 εd2 γd2 βd2 Wz2a bz2a Wz2b bz2b Wp2 bp2 εp2 γp2 βp2 mbNoExpFwdB N Wd1 bd1 εd1 γd1 βd1 Wz1a bz1a Wz1b bz1b Wp1 bp1 εp1 γp1 βp1 stemB N Ws bs εs γs βs)

                                  The representative batched EfficientNet-B0 has a (correct) VJP. Chained from the per-block gradients — stem → MBConv1 → MBConv6-strided → MBConv6-residual → head — via vjp_comp. The backward is genuinely composed from the proven per-op VJPs, batchMap-lifted to the batch and crossing the batch-coupled true batch-norm; HasVJP.correct pins it to the true Jacobian-transpose. The batched, true-batch-norm + SE analogue of efficientnet_has_vjp (EfficientNet.lean).

                                  Stated on the -composition of the blocks — which IS efficientnetForwardB (its nested-application spelling, used by the forward proof, is definitionally this composition); vjp_comp builds exactly this composition, so the proof closes structurally without re-reducing the whole net.

                                  Equations
                                  • One or more equations did not get rendered due to their size.
                                  Instances For
                                    theorem Proofs.efficientnetForwardB_eq_chain (N : ) (Ws : Kernel4 32 3 3 3) (bs : Vec 32) (εs : ) (γs βs : Vec 32) (Wd1 : DepthwiseKernel 32 3 3) (bd1 : Vec 32) (εd1 : ) (γd1 βd1 : Vec 32) (Wz1a : Mat 32 8) (bz1a : Vec 8) (Wz1b : Mat 8 32) (bz1b : Vec 32) (Wp1 : Kernel4 16 32 1 1) (bp1 : Vec 16) (εp1 : ) (γp1 βp1 : Vec 16) (We2 : Kernel4 96 16 1 1) (be2 : Vec 96) (εe2 : ) (γe2 βe2 : Vec 96) (Wd2 : DepthwiseKernel 96 3 3) (bd2 : Vec 96) (εd2 : ) (γd2 βd2 : Vec 96) (Wz2a : Mat 96 4) (bz2a : Vec 4) (Wz2b : Mat 4 96) (bz2b : Vec 96) (Wp2 : Kernel4 24 96 1 1) (bp2 : Vec 24) (εp2 : ) (γp2 βp2 : Vec 24) (We3 : Kernel4 144 24 1 1) (be3 : Vec 144) (εe3 : ) (γe3 βe3 : Vec 144) (Wd3 : DepthwiseKernel 144 5 5) (bd3 : Vec 144) (εd3 : ) (γd3 βd3 : Vec 144) (Wz3a : Mat 144 6) (bz3a : Vec 6) (Wz3b : Mat 6 144) (bz3b : Vec 144) (Wp3 : Kernel4 24 144 1 1) (bp3 : Vec 24) (εp3 : ) (γp3 βp3 : Vec 24) (Wh : Kernel4 1280 24 1 1) (bh : Vec 1280) (εh : ) (γh βh : Vec 1280) (Wfc : Mat 1280 10) (bfc : Vec 10) (x : Vec (N * (3 * 224 * 224))) :
                                    efficientnetForwardB N Ws bs εs γs βs Wd1 bd1 εd1 γd1 βd1 Wz1a bz1a Wz1b bz1b Wp1 bp1 εp1 γp1 βp1 We2 be2 εe2 γe2 βe2 Wd2 bd2 εd2 γd2 βd2 Wz2a bz2a Wz2b bz2b Wp2 bp2 εp2 γp2 βp2 We3 be3 εe3 γe3 βe3 Wd3 bd3 εd3 γd3 βd3 Wz3a bz3a Wz3b bz3b Wp3 bp3 εp3 γp3 βp3 Wh bh εh γh βh Wfc bfc x = (headFwdB N Wh bh εh γh βh Wfc bfc mbResidFwdB N We3 be3 εe3 γe3 βe3 Wd3 bd3 εd3 γd3 βd3 Wz3a bz3a Wz3b bz3b Wp3 bp3 εp3 γp3 βp3 mbStridedFwdB N We2 be2 εe2 γe2 βe2 Wd2 bd2 εd2 γd2 βd2 Wz2a bz2a Wz2b bz2b Wp2 bp2 εp2 γp2 βp2 mbNoExpFwdB N Wd1 bd1 εd1 γd1 βd1 Wz1a bz1a Wz1b bz1b Wp1 bp1 εp1 γp1 βp1 stemB N Ws bs εs γs βs) x

                                    ⭐⭐ THE SHAPE CHECK — efficientnetForwardB IS the -chain efficientnetForwardB_has_vjp is stated on. The committed forward is written in nested-application form so its faithfulness proof closes by pure delta; the apex above is stated on the composition vjp_comp builds. Until now the bridge between the two lived in that apex's docstring — "its nested-application spelling … is definitionally this composition".

                                    A justification in a docstring is what stops anyone re-checking it. Every whole-net tie takes its blocks OPAQUE so the isDefEq compares variables, which means the tie's subject is a chain of variables and nothing in it says which net they are — that is how ResNet-34's tie reversed the 2×2 pool for a month while the committed forward pooled 3×3/s2, and it moved a committed number 4×. The peers are mobilenetv2Forward_full_pc_eq_chain, convNextForwardTCh_eq_chain, resnet34Forward_full_pc_eq_chain and, for the 16-block net, efficientnetForwardB_full_eq_chain — this is the same theorem for the 3-block batched representative, which is the net efficientnetInputGradB reverses.

                                    ⚠ PROOF SHAPE MATTERS, and efficientnetForwardB_full_eq_chain records why: the equation-lemma rw plus one Function.comp_apply per stage closes syntactically, where a simp/rfl of the same statement makes the kernel reduce the block bodies and time out.