The decimation index map: a small output flat index k ↔ (co, ho, wo) maps to the even input position (co, 2·ho, 2·wo) in the (2h)×(2w) grid. A pure reindex Fin (oc·h·w) → Fin (oc·2h·2w); decimateFlat reads through it.

Flat spatial decimation Vec (oc·2h·2w) → Vec (oc·h·w): keep the even spatial positions. A coordinate reindex fun y k => y (decimateIdx k) — i.e. reindexCLM decimateIdx — so it is continuous-linear (hence differentiable), and decimate2 (conv2d …) = conv_stride2 ….

Decimation VJP. decimateFlat is a reindex, so its Jacobian is the sparse δ(idx = decimateIdx j) (pdiv_reindex); the backward scatters dy back to the even positions (zero elsewhere) — the "zero-upsampling" that StableHLO renders as lhs_dilation = [2,2]. reindexVJP at the decimation index.

Decimate after any VJP-carrying map: vjpComp with decimateFlatHasVJP. Every strided op below is decimateFlat ∘ (its stride-1 op), so each strided VJP is this applied to the stride-1 one; the expected type unfolds the strided op's definition.

Stride-2 SAME convolution, flattened: Vec (ic·2h·2w) → Vec (oc·h·w). Defined as decimateFlat ∘ flatConv (the stride-1 SAME conv on the 2h×2w grid, then keep even positions) — provably the genuine stride-2 conv.

Stride-2 conv input-VJP — the centerpiece. By the chain rule (vjpComp) on decimateFlat ∘ flatConv, reusing the proven stride-1 conv input-VJP (conv2dHasVJP3 via the flatten bridge) and the decimation VJP. The backward is flatConv.back (decimate.back dy) — i.e. zero-upsample the cotangent, then run the reversed-kernel conv (StableHLO: lhs_dilation=[2,2] on the transpose-reverse convolution).

Stride-2 conv weight-VJP. The same composition decimate ∘ conv viewed as a function of the kernel (input x fixed): the weight-grad is conv_weight_grad run on the zero-upsampled cotangent. By vjpComp, reusing the proven stride-1 conv2dWeightGradHasVJP + decimateFlatHasVJP.

Stride-2 conv bias-VJP. fun b => flatConvStride2 W b x = decimate ∘ (conv2d-in-b); by vjpComp of the proven stride-1 conv2dBiasGradHasVJP with decimateFlatHasVJP. The bias peer of flatConvStride2WeightGradHasVJP.

The odd decimation index map: like decimateIdx but keeping the odd positions (co, 2·ho+1, 2·wo+1). Composed under an even decimation it reads 4·ho+1 — which for the SAME conv at pad (k-1)/2 = 1 (k = 4) makes the window exactly the left-aligned x[4i .. 4i+3]: the real (paper/render) pad-0 stride-4 patchify, never touching the boundary padding.

Flat odd spatial decimation Vec (oc·2h·2w) → Vec (oc·h·w): keep the odd spatial positions. A coordinate reindex, exactly as decimateFlat.

Odd-decimation VJP — the same sparse-δ reindex Jacobian as decimateFlatHasVJP, at the odd positions.

Odd-decimate after any VJP-carrying map: the odd-phase peer of HasVJP.decimate.

Stride-4 patchify convolution, flattened: Vec (ic·4h·4w) → Vec (oc·h·w). decimateFlat ∘ decimateOddFlat ∘ (stride-1 SAME conv) — reads the SAME conv (pad (k-1)/2) at positions 4i+1, which for the 4×4 stem is the left-aligned window x[4i .. 4i+3]: the paper's pad-0 Conv2d(k=4, s=4) and the committed render's patchify, in-bounds at every tap.

Stride-4 conv input-VJP — two vjpComp steps over the proven stride-1 conv input-VJP and the two decimation VJPs (backward = zero-upsample twice, then the reversed-kernel conv).

Stride-4 conv weight-VJP. The kernel-side peer of flatConvStride4HasVJP, and the stride-4 analogue of flatConvStride2WeightGradHasVJP: the same decimateFlat ∘ decimateOddFlat ∘ conv composition viewed as a function of the kernel (input x fixed), so the weight-grad is conv2dWeightGrad run on the twice-zero-upsampled cotangent. Two vjpComp steps over the proven stride-1 weight-VJP and the two decimation VJPs — no new mathematics, only the composition the stride-2 sibling already does once.

This is the weight-gradient cert of ConvNeXt's 4×4/s4 patchify stem (psW); its forward (flatConvStride4) and input-VJP (flatConvStride4HasVJP) are proven above.

Stride-2 XLA-SAME convolution, flattened: Vec (ic·2h·2w) → Vec (oc·h·w). decimateOddFlat ∘ flatConv — the stride-1 symmetric-SAME conv on the 2h×2w grid, then keep the odd positions. The asymmetric-pad peer of flatConvStride2.

Stride-2 XLA-SAME input-VJP. vjpComp on decimateOddFlat ∘ flatConv, reusing the proven stride-1 conv input-VJP and the odd-decimation VJP. The backward zero-upsamples the cotangent onto the odd positions and then runs the reversed-kernel conv — i.e. StableHLO's lhs_dilation = [2,2] with the transposed padding shifted by one, which is exactly the asymmetry the forward introduced. A symmetric backward against this forward is a silent wrong-gradient, and it is this composition that rules it out: the offset lives in one place and both directions read it.

Stride-2 XLA-SAME weight-VJP. The same composition viewed as a function of the kernel (input x fixed): conv2dWeightGrad run on the odd-zero-upsampled cotangent.

Stride-2 XLA-SAME bias-VJP. fun b => flatConvStride2Xla W b x = the odd decimation of the stride-1 conv-in-b; by vjpComp of conv2dBiasGradHasVJP with the odd-decimation VJP. The bias peer of flatConvStride2BiasGradHasVJP.