maxPool3s2 — the 3×3 stride-2 max pool of He et al.'s ResNet stem #
Written 2026-08-03. planning/archive/rsb_a3_r50_verified.md / the stem-pool finding.
The verified path's only pooling op is maxPool2 — 2×2, stride 2, non-overlapping. Every
ResNet in He et al. (18/34/50/101/152) specifies a 3×3 stride-2 pool after the stem conv, so
resnet34Verified has been pooling a different function from the paper and from the reference
it is paired against. This file is the missing op.
⚠⚠ WHICH 3×3 pool — the PAPER's, which is not what the JAX reference emits today #
This file implements He et al. / torchvision: nn.MaxPool2d(3, stride=2, padding=1) —
symmetric padding, so window i covers input [2i−1, 2i+1].
⚠ The repo's JAX references emit reduce_window(…, (1,1,3,3), (1,1,2,2), 'SAME'), and XLA's
SAME on a 112→56 axis gives pad_total = max((56−1)·2 + 3 − 112, 0) = 1, split
pad_low = 0, pad_high = 1 — window i = [2i, 2i+2], padded at the end. Measured on device
at n = 12: SAME windows peak at [2,4,6,8,10,11], symmetric at [1,3,5,7,9,11]. The two
grids are offset by one input position and are different functions everywhere.
Paper-faithfulness is the goal, so this file is symmetric and the JAX max_pool2d helper moves to
match. ⚠ Until both land and are re-run, verified and JAX disagree at the stem pool.
⭐ Measured: symmetric (k−1)//2 padding is bit-identical to SAME for every 2×2 pool, so no
cifar/mnist net moves — only the 3×3 users (R34-ImageNet, R50).
⭐ THE PADDING NEEDS NO EXTENDED-REALS TYPE #
reduce_window pads with -∞, which Tensor3 _ _ _ = … → ℝ cannot hold. It does not need to:
for max, clamping the index is equivalent to -∞ padding. The only out-of-range read is
2i−1 at i = 0, and Nat's truncated subtraction clamps it to 0 — a cell the window already
contains at offset a = 1. So max over the clamped triple equals max over the unpadded pair,
which is exactly what -∞ padding computes. win3RowInv_first_dup is that statement.
⭐ The symmetric form needs no min: the upper end 2(h−1)+2−1 = 2h−1 is in range by
construction, so truncated subtraction is the whole story.
The shape of the VJP, and why overlap costs less than it looks #
maxPool2's backward is a lookup (dy at hi/2), sound only because 2×2 windows are
disjoint. Here windows overlap — odd input p lies in windows (p−1)/2 and (p+1)/2 — so an
input feeds up to 4 outputs and the backward must accumulate.
That needs no new analytic argument. HasVJPAt3.correct already states the backward as
∑ co ∑ ho ∑ wo, pdiv3 f x … * dy co ho wo — a sum over all outputs; maxPool2's
codegen_matches_canonical merely collapses it to one term using disjointness. Here it
collapses to ≤4. The generic route (maxPool2LocalReindex → reindexCLM → pdiv3) is
indifferent: at a smooth point the pool is locally a reindexing map, and overlap only makes that
map non-injective, which reindexCLM's adjoint already handles by summing over preimages.
⚠ This file is the DEFINITIONAL layer. The HasVJPAt3 witness mirrors
maxPool2_has_vjp_at3 (33 declarations in CNN.lean) and is not built here yet. Nothing
downstream references these definitions, so nothing renders or moves.
Input row of offset a ∈ Fin 3 inside output window hi_out: 2·hi + a − 1, in Nat.
The truncated subtraction is the low pad (see the header) — it is the only clamp needed,
because the high end 2(h−1)+2−1 = 2h−1 is in range by construction.
Instances For
Column peer of win3RowInv.
Instances For
⭐ The padding statement. In the FIRST window offset a = 0 duplicates a = 1 rather
than reading out of range — exactly what a -∞ pad contributes to a max.
3×3 stride-2 symmetrically-padded max pool, [c, 2h, 2w] → [c, h, w]: the max over the
window [2i−1, 2i+1] × [2j−1, 2j+1], clamped at the near edge (= -∞ padded, header).
Equations
- Proofs.maxPool3s2 x ch hi wi = Finset.univ.sup' Proofs.maxPool3s2._proof_1 fun (ab : Fin 3 × Fin 3) => x ch (Proofs.win3RowInv hi ab.1) (Proofs.win3ColInv wi ab.2)
Instances For
The pooled value is attained by some window cell.
The 3×3 pool never grows magnitudes — it selects an existing window cell.
⭐ Finset.sup' again makes the window size stop mattering: maxPool2_abs_le needs a nested
abs_max_le (abs_max_le _ _) (abs_max_le _ _) for 4 cells, which at 9 would be worse; here it
is sup'_le plus one le_sup', independent of the window. Fourth collapse of the same kind.
The 3×3 pool is 1-Lipschitz in the sup norm — the peer of maxPool2_close. Standard
sup-vs-sup argument: each side is ≤ the other plus e, from sup'_le and le_sup'.
Smoothness: every 3×3 window has a strict argmax. Stated as "distinct offsets that land
on distinct input POSITIONS have distinct values", so the clamped duplicate in the first
window (a = 0 ≡ a = 1) is not counted as a tie. ⚠ That carve-out is forced by the padding
and has no maxPool2 analogue — there, distinct offsets always meant distinct positions.
Equations
- One or more equations did not get rendered due to their size.
Instances For
⭐ Positional injectivity ⇒ MaxPool3s2Smooth — the discharge lemma for the whole-net live
and seal witnesses, the peer of MnistCNN's maxPool2Smooth_of_injective. One injectivity
argument in place of 36·c·h·w per-window decides (9 offsets pairwise, against 2×2's 6), which
at ResNet-34's stem is why case-bashing is not an option.
⚠⚠ And it is STRICTLY SHORTER than its 2×2 peer, for the reason the padding forced.
MaxPool2Smooth is quantified over offsets, so its discharge lemma has to get from
"the two input positions coincide" back to "the two offsets coincide" — two Fin.mk.injEq +
omega decodes, valid only because there distinct offsets always meant distinct positions.
MaxPool3s2Smooth is quantified over positions precisely because that is false here (the
clamped duplicate a = 0 ≡ a = 1 in the first window, win3RowInv_first_dup), so injectivity
lands directly on the hypothesis and the decode step does not exist. The carve-out that made
the predicate awkward to state is what makes it cheap to discharge.
The hypothesis is the same one the 2×2 sites already supply: on each channel the position map
(r, s) ↦ x ci r s is injective.
Input (ci, hi_in, wi_in) attains the max of the window at output (ho, wo).
⚠ Unlike MaxPool2IsArgmax this takes the OUTPUT position explicitly: with overlapping
windows an input has no single owning window, which is the whole difference.
Equations
- One or more equations did not get rendered due to their size.
Instances For
⭐ The overlap fact, stated rather than assumed: an input row lies in at most TWO windows —
p/2 and (p+1)/2. With symmetric padding the shared cell is at ODD p (window (p−1)/2
takes it at offset 2, window (p+1)/2 at offset 0); even p lies in exactly one. So an input
feeds at most 4 outputs and the backward accumulates at most 4 terms — the count maxPool2
does not have.
A (not necessarily unique) argmax of the 3×3 window at output (co, ho, wo).
Unique under MaxPool3s2Smooth up to position — the clamped duplicate in the first
window is two offsets naming one cell.
Equations
- Proofs.maxPool3s2Argmax x co ho wo = Classical.choose ⋯
Instances For
If (a, b) dominates every window cell, the pooled value is the value there.
⭐ The sup' formulation makes this two lines where maxPool2_eq_at_max needs a
four-way fin_cases against an explicit max (max _ _) (max _ _) — and nine ways here.
For each output flat index, the flat index of its argmax's input position.
⚠ Not injective — two overlapping windows may select the same input. That is exactly
what makes the backward accumulate, and reindexCLM's adjoint already sums over preimages,
so nothing here needs to change relative to maxPool2LocalReindex.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Smooth-point local linearisation. On a metric ball around flatten x the flattened
pool agrees with the reindex y ↦ y ∘ σ. Mirrors maxPool2_flat_hasFDerivAt; the one
structural change is that the gap function and the domination argument branch on whether the
two offsets name the same position (the clamped duplicate), not on whether the offsets
are equal — MaxPool3s2Smooth says nothing about coincident positions because there the
values are literally the same number.
Smooth-point Jacobian. pdiv3 is the 0/1 indicator that the local reindex sends output
(co, ho, wo) to input (ci, hi_in, wi_in).
⚠⚠ This is where the overlapping case genuinely differs from maxPool2, and it differs by
being SIMPLER to state. pdiv3_maxPool2_smooth decodes the condition into
co = ci ∧ ho = winRow hi_in ∧ wo = winCol wi_in ∧ IsArgmax — legitimate there because each
input has exactly ONE owning window, so winRow/winCol name it. Here an input has up to
two windows per axis and no such decoding exists. Leaving the condition as the reindex
equation is both correct and shorter; the accumulation then happens in correct's sum over
outputs, with no extra argument.
The VJP witness. The backward accumulates dy over every output whose window selects
this input — at most 4 of them (win3Row_mem_le_two squared). maxPool2's peer is a single
lookup; this is the same statement without the disjointness collapse.
Equations
- One or more equations did not get rendered due to their size.
Instances For
Flattened 3×3/s2 pool, the Vec-level form the codegen denotes.
Equations
- Proofs.maxPool3s2Flat c h w v = (Proofs.maxPool3s2 (Proofs.Tensor3.unflatten v)).flatten
Instances For
Equations
- Proofs.maxPool3s2Flat_has_vjp_at x h_smooth = Proofs.hasVJPAt3_to_hasVJPAt (Proofs.maxPool3s2_has_vjp_at3 x h_smooth)