Batched activation stages — conv/depthwise → true BN → relu or relu6, as CertLayers #
The kinked stages the ResNets, MobileNetV2 and MobileNetV4 are built from, at the batched index
N·(c·h·w). Each comes with its _at VJP (certified where the pre-activation misses the kink),
its backward graph (.selectPos for relu, .selectMid for relu6) with _faithful, and a
CertLayer bundling the two. The smooth stages (swish, projection) are in BatchedStages.
| stage | forward | CertLayer |
|---|---|---|
conv → BN → relu6; depthwise → BN → relu6 (stride 1 / XLA-SAME stride 2) | cbrB, dwbrB, dwbrBstrided | cbrLayer, dwbrLayer, dwbrStridedLayer |
| 1×1 projection → BN | projB | projLayer |
| conv → BN → relu; strided conv → BN → relu | cbReluB, cbReluStridedB | cbReluLayer, cbReluStridedLayer |
| strided projection → BN (the downsample skip) | projStridedB | projStridedLayer |
Batched conv → bn → relu6 stage (MobileNetV2 expand), at the network layout
N·(oc·h·w). Identical to EfficientNet's cbsB but with relu6 for swish.
Equations
- Proofs.StableHLO.cbrB N W b ε γ β = Proofs.relu6 (N * (oc * h * w)) ∘ Proofs.StableHLO.bnBatchLA N oc h w ε γ β ∘ Proofs.StableHLO.batchMap N (Proofs.flatConv W b)
Instances For
Batched depthwise → bn → relu6 stage (MobileNetV2 depthwise), at the network
layout. Identical to EfficientNet's dwbsB but with relu6 for swish.
Equations
- Proofs.StableHLO.dwbrB N W b ε γ β = Proofs.relu6 (N * (c * h * w)) ∘ Proofs.StableHLO.bnBatchLA N c h w ε γ β ∘ Proofs.StableHLO.batchMap N (Proofs.depthwiseFlat W b)
Instances For
Batched STRIDE-2 depthwise → bn → relu6 stage (MobileNetV2 downsample
depthwise), at the network layout. The stride-2 analogue of dwbrB: maps the
larger input spatial c·(2h)·(2w) to the output spatial c·h·w. Identical to
EfficientNet's dwbsSB but with relu6 for swish.
Equations
- Proofs.StableHLO.dwbrBstrided N W b ε γ β = Proofs.relu6 (N * (c * h * w)) ∘ Proofs.StableHLO.bnBatchLA N c h w ε γ β ∘ Proofs.StableHLO.batchMap N (Proofs.depthwiseStride2FlatXla W b)
Instances For
Generic relu6-on-batched-bn-stage _at VJP. The relu6 analogue of
bnSwishStage_has_vjp, but _at (relu6 only has a pointwise VJP): compose the
batched-op VJP, the true-BN VJP (both global, lifted via .toHasVJPAt), and
relu6's pointwise VJP at the pre-relu6 activation.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Differentiability of the generic relu6-on-batched-bn-stage at a smooth point.
cbrB (conv-bn-relu6) _at VJP at a smooth point.
Equations
- Proofs.StableHLO.cbrB_has_vjp_at N W b ε hε γ β x h_smooth = Proofs.StableHLO.bnRelu6Stage_has_vjp_at N (Proofs.flatConv W b) ⋯ (Proofs.flatConv_has_vjp W b) ε hε γ β x h_smooth
Instances For
dwbrB (depthwise-bn-relu6) _at VJP at a smooth point.
Equations
- Proofs.StableHLO.dwbrB_has_vjp_at N W b ε hε γ β x h_smooth = Proofs.StableHLO.bnRelu6Stage_has_vjp_at N (Proofs.depthwiseFlat W b) ⋯ (Proofs.depthwiseFlat_has_vjp W b) ε hε γ β x h_smooth
Instances For
dwbrBstrided (STRIDE-2 depthwise-bn-relu6) _at VJP at a smooth point. The
stride-2 analogue of dwbrB_has_vjp_at: lifts depthwiseStride2FlatXla_has_vjp
(the strided per-channel conv input-VJP) through the generic relu6-bn stage.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Batched conv → bn → relu6 stage backward graph (MobileNetV2 expand):
convBackBatched ∘ bnBatchLABack ∘ selectMid, each at its cumulative forward
activation. The relu6 analogue of cbsBackBatchedGraph — .selectMid (the
relu6 two-sided-kink mask) replaces .swishBack.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Batched depthwise → bn → relu6 stage backward graph (MobileNetV2 depthwise).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Batched STRIDE-2 depthwise → bn → relu6 stage backward graph (MobileNetV2
downsample depthwise). The stride-2 analogue of dwbrBackBatchedGraph: the
bn/relu6 run at the OUTPUT spatial h×w, then depthwiseStridedBackBatched maps
the bn-cotangent back to the larger input c·(2h)·(2w) (zero-upsample +
reversed-kernel per-channel depthwise). The relu6 back is .selectMid as before.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The conv → bn → relu6 stage as a CertLayer, certified where its pre-relu6 activation misses
both kinks. The kernel extent is a binder, so the same layer is a 1×1 or a 3×3.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The depthwise → bn → relu6 stage as a CertLayer.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The stride-2 depthwise → bn → relu6 stage as a CertLayer: 2h × 2w in, h × w out.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The conv → bn stage (projB, no activation) as a CertLayer. Globally certified
(ok = True): with no activation there is no kink.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Batched conv → bn → relu stage (ResNet basic-block first stage), at the
network layout N·(oc·h·w). The relu analogue of MobileNetV2's cbrB
(relu for relu6).
Equations
- Proofs.StableHLO.cbReluB N W b ε γ β = Proofs.relu (N * (oc * h * w)) ∘ Proofs.StableHLO.bnBatchLA N oc h w ε γ β ∘ Proofs.StableHLO.batchMap N (Proofs.flatConv W b)
Instances For
Generic relu-on-batched-bn-stage _at VJP. The relu analogue of
bnRelu6Stage_has_vjp_at (and of bnSwishStage_has_vjp, but _at — relu
only has a pointwise VJP): compose the batched-op VJP, the true-BN VJP (both
global, lifted via .toHasVJPAt), and relu's pointwise VJP at the pre-relu
activation. The smoothness hypothesis is the one-sided ≠ 0.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Differentiability of the generic relu-on-batched-bn-stage at a smooth point.
cbReluB (conv-bn-relu) _at VJP at a smooth point.
Equations
- Proofs.StableHLO.cbReluB_has_vjp_at N W b ε hε γ β x h_smooth = Proofs.StableHLO.bnReluStage_has_vjp_at N (Proofs.flatConv W b) ⋯ (Proofs.flatConv_has_vjp W b) ε hε γ β x h_smooth
Instances For
Batched conv → bn → relu stage backward graph (ResNet basic-block stage 1):
convBackBatched ∘ bnBatchLABack ∘ selectPos, each at its cumulative forward
activation. The relu analogue of MobileNetV2's cbrBackBatchedGraph —
.selectPos (the relu one-sided-kink mask) replaces .selectMid.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The conv → bn → relu stage as a CertLayer, certified where its pre-relu activation misses 0.
The kernel extent is a binder, so the same layer is a 1×1 or a 3×3.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Batched STRIDE-2 conv → bn → relu stage (downsample basic-block first
stage), at the network layout N·(oc·h·w) ← N·(ic·(2h)·(2w)). The strided
sibling of cbReluB (flatConvStride2 for flatConv); halves spatial.
Equations
- Proofs.StableHLO.cbReluStridedB N W b ε γ β = Proofs.relu (N * (oc * h * w)) ∘ Proofs.StableHLO.bnBatchLA N oc h w ε γ β ∘ Proofs.StableHLO.batchMap N (Proofs.flatConvStride2 W b)
Instances For
cbReluStridedB (strided conv-bn-relu) _at VJP at a smooth point. The strided
sibling of cbReluB_has_vjp_at (flatConvStride2 for flatConv).
Equations
- One or more equations did not get rendered due to their size.
Instances For
Batched strided conv → bn → relu stage backward graph:
convStridedBackBatched ∘ bnBatchLABack ∘ selectPos, each at its cumulative
forward activation. The strided sibling of cbReluBackBatchedGraph —
convStridedBackBatched (the new stride-2 batched-conv VJP) replaces
convBackBatched.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The strided conv → bn → relu stage as a CertLayer — cbReluLayer with flatConvStride2.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Batched strided conv → bn projection skip (downsample basic-block skip):
bnBatchLA ∘ batchMap (flatConvStride2) — the 3×3 stride-2 projection that
matches the body's downsampled oc·h·w output. The strided sibling of projB
(flatConvStride2 for flatConv); no activation (linear bottleneck).
Equations
- Proofs.StableHLO.projStridedB N W b ε γ β = Proofs.StableHLO.bnBatchLA N oc h w ε γ β ∘ Proofs.StableHLO.batchMap N (Proofs.flatConvStride2 W b)
Instances For
Equations
- Proofs.StableHLO.projStridedB_has_vjp N W b ε hε γ β = Proofs.bnStage_has_vjp N (Proofs.flatConvStride2 W b) ⋯ (Proofs.flatConvStride2_has_vjp W b) ε hε γ β
Instances For
Batched strided conv → bn projection-skip backward graph:
convStridedBackBatched ∘ bnBatchLABack, at the skip's forward activation. The
strided sibling of projBackBatchedGraph.
Equations
- One or more equations did not get rendered due to their size.
Instances For
The strided conv → bn projection skip as a CertLayer — projLayer with flatConvStride2,
globally certified.
Equations
- One or more equations did not get rendered due to their size.