Documentation

LeanMlir.Proofs.Foundation.Batched.Stages

Batched stages — conv/depthwise → true batch-norm (→ swish) at the batched index, with VJPs #

The building blocks every batched conv-net forward is written in, at the flat batched index N·(c·h·w): each stage is batchMap N of a proven per-example op, then true batch-norm (bnBatchLA, which couples the batch), then pointwise swish where the net has it.

stageforwardVJP
1×1 / stride-1 conv → BN → swishcbsBcbsBHasVJP
XLA-SAME stride-2 stem conv → BN → swishstemBstemBHasVJP
depthwise (stride 1 / symmetric stride 2) → BN → swishdwbsB / dwbsSBdwbsBHasVJP / dwbsSBHasVJP
squeeze-exciteseBseBHasVJP
1×1 projection → BN (no activation)projBprojBHasVJP

The generic pieces they compose — batchMapHasVJP (block-diagonal VJP of a batch-separable op) and bnBatchLAHasVJP — are in BatchMapVJPAt. EfficientNet-B0, MobileNetV2/V4 and the ResNets build their blocks from these; the MBConv blocks themselves are in EfficientNetRender.PC and EfficientNetChainClose.

source
@[reducible]
noncomputable def Proofs.cbsB (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (γ β : Vec oc) :
Vec (N * (ic * h * w)) → Vec (N * (oc * h * w))

Batched conv → bn → swish (1×1 expand / generic stride-1 conv).

Equations
Instances For
    source
    noncomputable def Proofs.stemB (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (γ β : Vec oc) :
    Vec (N * (ic * (2 * h) * (2 * w))) → Vec (N * (oc * h * w))

    Batched strided (3×3 s2) stem conv → bn → swish (halves spatial). Note: at the XLA-SAME phase (flatConvStride2Xla = decimateOddFlat ∘ flatConv): the TF-origin B0 pads its stem (0,1), and the shipped render emits convStridedXla there. The symmetric flatConvStride2 has the same type and output shape; nothing structural would notice the wrong one.

    Equations
    Instances For
      source
      @[reducible]
      noncomputable def Proofs.dwbsB (N : ℕ) {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ℝ) (γ β : Vec c) :
      Vec (N * (c * h * w)) → Vec (N * (c * h * w))

      Batched depthwise (stride-1, k×k) → bn → swish.

      Equations
      Instances For
        source
        @[reducible]
        noncomputable def Proofs.dwbsSB (N : ℕ) {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ℝ) (γ β : Vec c) :
        Vec (N * (c * (2 * h) * (2 * w))) → Vec (N * (c * h * w))

        Batched depthwise (stride-2 downsample, k×k) → bn → swish.

        Equations
        Instances For
          source
          @[reducible]
          noncomputable def Proofs.seB (N : ℕ) {c h w r : ℕ} (W₁ : Mat c r) (b₁ : Vec r) (W₂ : Mat r c) (b₂ : Vec c) :
          Vec (N * (c * h * w)) → Vec (N * (c * h * w))

          Batched squeeze-excite block x ⊙ gate(x) (the proven seBlockFull, per example).

          Equations
          Instances For
            source
            @[reducible]
            noncomputable def Proofs.projB (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (γ β : Vec oc) :
            Vec (N * (ic * h * w)) → Vec (N * (oc * h * w))

            Batched project: 1×1 conv → bn (no swish — the linear bottleneck).

            Equations
            Instances For
              source
              noncomputable def Proofs.flatConvHasVJP {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) :
              HasVJP (flatConv W b)

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

              Equations
              Instances For
                source
                theorem Proofs.bnSwishStage_differentiable (N : ℕ) {a oc h w : ℕ} (op : Vec a → Vec (oc * h * w)) (hop : Differentiable ℝ op) (ε : ℝ) (hε : 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.

                source
                noncomputable def Proofs.bnSwishStageHasVJP (N : ℕ) {a oc h w : ℕ} (op : Vec a → Vec (oc * h * w)) (hop : Differentiable ℝ op) (hopv : HasVJP op) (ε : ℝ) (hε : 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
                  source
                  theorem Proofs.bnStage_differentiable (N : ℕ) {a oc h w : ℕ} (op : Vec a → Vec (oc * h * w)) (hop : Differentiable ℝ op) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :
                  Differentiable ℝ (StableHLO.bnBatchLA N oc h w ε γ β ∘ StableHLO.batchMap N op)

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

                  source
                  noncomputable def Proofs.bnStageHasVJP (N : ℕ) {a oc h w : ℕ} (op : Vec a → Vec (oc * h * w)) (hop : Differentiable ℝ op) (hopv : HasVJP op) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :
                  HasVJP (StableHLO.bnBatchLA N oc h w ε γ β ∘ StableHLO.batchMap N op)

                  VJP of a batched conv → bn stage.

                  Equations
                  • One or more equations did not get rendered due to their size.
                  Instances For
                    source
                    theorem Proofs.cbsB_differentiable (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :
                    Differentiable ℝ (cbsB N W b ε γ β)
                    source
                    noncomputable def Proofs.cbsBHasVJP (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :
                    HasVJP (cbsB N W b ε γ β)
                    Equations
                    Instances For
                      source
                      theorem Proofs.stemB_differentiable (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :
                      Differentiable ℝ (stemB N W b ε γ β)
                      source
                      noncomputable def Proofs.stemBHasVJP (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :
                      HasVJP (stemB N W b ε γ β)
                      Equations
                      Instances For
                        source
                        theorem Proofs.dwbsB_differentiable (N : ℕ) {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ℝ) (hε : 0 < ε) (γ β : Vec c) :
                        Differentiable ℝ (dwbsB N W b ε γ β)
                        source
                        noncomputable def Proofs.dwbsBHasVJP (N : ℕ) {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ℝ) (hε : 0 < ε) (γ β : Vec c) :
                        HasVJP (dwbsB N W b ε γ β)
                        Equations
                        Instances For
                          source
                          theorem Proofs.dwbsSB_differentiable (N : ℕ) {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ℝ) (hε : 0 < ε) (γ β : Vec c) :
                          Differentiable ℝ (dwbsSB N W b ε γ β)
                          source
                          noncomputable def Proofs.dwbsSBHasVJP (N : ℕ) {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (ε : ℝ) (hε : 0 < ε) (γ β : Vec c) :
                          HasVJP (dwbsSB N W b ε γ β)
                          Equations
                          Instances For
                            source
                            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₂)
                            source
                            noncomputable def Proofs.seBHasVJP (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
                              source
                              theorem Proofs.projB_differentiable (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :
                              Differentiable ℝ (projB N W b ε γ β)
                              source
                              noncomputable def Proofs.projBHasVJP (N : ℕ) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :
                              HasVJP (projB N W b ε γ β)
                              Equations
                              Instances For