A conv kernel: out_channels × in_channels × kH × kW. This is the OIHW layout used by StableHLO and IREE.

Row-major flatten: Kernel4 oc ic kH kW → Vec (oc * ic * kH * kW).

Row-major unflatten: inverse of flatten.

Conv2d forward (SAME padding, stride 1).

y[o, h, w] = (Σ_{c, kh, kw} x[c, h+kh−p, w+kw−p] · W[o, c, kh, kw]) + b[o]

where p = (kH−1)/2 is the padding offset and out-of-bounds reads return 0 (zero padding). The output spatial size equals the input.

Note: this is technically cross-correlation, not convolution in the strict signal-processing sense. ML literature uses "convolution" loosely; the difference (kernel flipping) only matters when comparing against classical signal-processing references.

MLIR (hand_cnn_train_step.mlir): %cv0 = "stablehlo.convolution"(%x, %W0) { padding = dense<[[1, 1], [1, 1]]>, ... } %h0pre = stablehlo.add %cv0, broadcast(%b0)

Closed-form input gradient for conv2d — direct formula, written as a sum over output positions (co, ho, wo) with reconstructed kernel offsets kh_nat = hi + pH − ho, kw_nat = wi + pW − wo. The body is nonzero only when the reconstructed (kh_nat, kw_nat) lies in [0, kH) × [0, kW) — i.e., when the input position (hi, wi) is actually reachable from output (ho, wo) via some valid kernel offset. Equivalent (under the (ho, wo) ↔ (kh, kw) partial bijection ho = hi+pH-kh) to the MLIR-aligned "reversed-kernel" formula dx[c, h, w] = Σ_{o, kh, kw} W[o, c, kH−1−kh, kW−1−kw] · dy[o, h+kh−p, w+kw−p].

Conv2d input-VJP — proved from foundation rules.

The function v ↦ flatten (conv2d W b (unflatten v)) is affine in v: the bias-free conv (linear in v) plus the broadcast bias, so pdiv_of_affine reads each per-(idx_in, idx_out) entry off the bias-free conv of the basis vector — a sum over (c, kh, kw) of W o(idx_out) c kh kw times the pad-guarded Kronecker. Reindex Fin (oc*h*w) ↔ Fin oc × Fin h × Fin w on the sum-over-idx_out; padTap_indicator turns each Kronecker into "c = ci and the tap lands on (hi, wi)", and sum_fin_ite_add_eq collapses (c, kh, kw) to the closed-form input gradient conv2dInputGradFormula.

The backward function (accessed as (conv2dHasVJP3 W b).backward, or via the conv2dInputGrad abbrev below) implements conv2dInputGradFormula W dy ci hi wi — a direct sum over (co, kh, kw) of W co ci kh kw * dy co ho_nat wo_nat for valid (ho_nat, wo_nat). Equivalent (under kh ↔ kH−1−kh) to the MLIR-aligned reversed-kernel formula dx[c, h, w] = Σ_{o, kh, kw} W[o, c, kH−1−kh, kW−1−kw] · dy[o, h+kh−p, w+kw−p].

MLIR emits the reversed-kernel form directly: %W1_t = stablehlo.transpose %W1, dims = [1, 0, 2, 3] -- swap oc↔ic %W1_rev = stablehlo.reverse %W1_t, dims = [2, 3] -- flip spatial %d_h0 = "stablehlo.convolution"(%d_h1pre, %W1_rev) ...

Named accessor for the conv2d input backward — aligns with MLIR codegen (stablehlo.convolution in the backward pass).

Flat conv — conv2d bridged into flattened Vec → Vec space: flatConv W b = flatten ∘ conv2d W b ∘ unflatten. Spatial dims are preserved (ic h w → oc h w), so this is Vec (ic*h*w) → Vec (oc*h*w). This is the form the ResNet/CNN VJP composition actually uses (everything lives in flat Vec space).

An act ∘ norm ∘ lin stage's VJP at a point — a conv (or depthwise, or batched) op, a normalisation, then an activation. lin and norm are differentiable everywhere, so their global VJPs lift through .toHasVJPAt; act needs a derivative and a VJP only at the stage's pre-activation norm (lin v), which is what relu and relu6 have off their kinks. Two vjpCompAts: norm ∘ lin, then act. Every conv-bn-act stage VJP in the repo is this at a particular lin, norm and act.

conv → bn → relu block VJP at a smooth point.

The workhorse for composing a ResNet VJP. In flattened Vec space, the block is relu ∘ bnForward ∘ flatConv : Vec (ic*h*w) → Vec (oc*h*w) (BatchNorm runs over the oc*h*w flattened activations with scalar ε, γ, β). It is stageHasVJPAt at a point v:

• inner = bnForward ∘ flatConv — both differentiable everywhere (flatConv_differentiable, bnForward_differentiable), so their bundled VJPs lift through .toHasVJPAt. The conv witness is the HasVJP3-bridged HasVJP3.toHasVJP (conv2dHasVJP3 W b).
• outer = relu — needs the smoothness hypothesis h_smooth (no post-BN activation hits the ReLU kink) for both relu_differentiableAt_of_smooth and reluHasVJPAt.

Mirrors mlpHasVJPAt (dense→relu→dense), with flatConv/bn in place of the dense layers.

conv → bn block VJP (no ReLU), everywhere. Just flatConv then bnForward, both differentiable everywhere, so this is a global HasVJP (no smoothness needed). This is the building block for the second conv→bn of a residual body and for the 1×1 projection skip.

Basic-block body VJP at a smooth point.

F := convBn₂ ∘ convBnRelu₁ = bn₂ ∘ conv₂ ∘ relu ∘ bn₁ ∘ conv₁, the body of a post-activation ResNet basic block (the outer ReLU and skip-add are applied later). Channels go ic → mid → oc (generic; set ic = mid = oc = c for the identity-skip block, ic ≠ oc for the downsample/projection block). Spatial dims h w preserved.

Inner convBnRelu₁ needs the smoothness hyp h_smooth₁ (no post-bn₁ activation hits the ReLU kink); outer convBn₂ is everywhere differentiable, lifted via .toHasVJPAt. Two vjpCompAt chain.

Full basic residual block VJP (identity skip).

relu(x + F(x)) with F the conv→bn→relu→conv→bn body and an identity skip (so ic = mid = oc = c, spatial preserved). Two smoothness hyps: h_smooth₁ for the inner block ReLU, and h_smooth_out for the post-add outer ReLU (F v + v avoids the kink). Built as relu ∘ residual F via residualHasVJPAt then a final vjpCompAt with relu.

Downsample/projection basic residual block VJP.

relu(proj(x) + F(x)) where the channel/stride change is folded into the conv dims: body F maps ic → oc (first conv ic → oc, second oc → oc), and the skip is a 1×1 convBn projection proj : ic → oc (everywhere differentiable — no ReLU, so its diffAt is immediate). Built with residualProjHasVJPAt, then vjpCompAt with the post-add relu under h_smooth_out.

Conv2d weight-VJP — proved from foundation rules.

The function v ↦ flatten (conv2d (unflatten v) b x) is affine in v: the broadcast bias b o(idx_out) plus the bias-free conv, which is linear in v. pdiv_of_affine reads each Jacobian entry off the bias-free conv of a basis vector — a Kronecker at (o', c', kh', kw') against the padded input — and the collapsed sum is the transpose-trick backward dW[o', c', kh', kw'] = Σ_{hi, wi} x_pad_term(...) · dy(flat(o', hi, wi)).

Named accessor for the conv2d weight backward — aligns with MLIR codegen (the "transpose trick" stablehlo.convolution in the backward pass). Unwraps the flattening so callers see Kernel4 → Kernel4.

Conv2d bias-VJP — proved from foundation rules. The function b ↦ flatten (conv2d W b x) decomposes as (channel broadcast of b) + (bias-free conv, constant in b), so pdiv_of_affine gives the channel Kronecker, collapsed over the (c, hi, wi) decomposition of Fin (oc*h*w). The backward is db[o] = Σ_{hi, wi} dy[o, hi, wi] (the same sum as conv2dBiasGradFormula below, written on the flattened dy).

Named accessor for the conv2d bias backward via the VJP framework.

Conv2d bias gradient — closed-form formula.

db[o] = Σ_{h, w} dy[o, h, w]

This is the sum conv2dBiasGradHasVJP's backward is defined as, on Tensor3 instead of the flattened dy. No theorem in this file states the equation between the two.

Each output cell adds the same b[o], so its gradient accumulates the contributions from every spatial position. MLIR emits this as a stablehlo.reduce across the spatial (and batch) dims.

MaxPool 2×2 stride 2 forward — concrete definition.

Each output cell is the maximum of a 2×2 window of input cells: y[c, h, w] = max{ x[c, 2h+a, 2w+b] : a, b ∈ {0,1} }, written as an explicit four-way max.

MLIR: %pool = "stablehlo.reduce_window"(%h1, %neginf) ({ ^bb0(%a, %b): stablehlo.return (stablehlo.maximum %a, %b) }) {window_dimensions = [1, 1, 2, 2], window_strides = [1, 1, 2, 2]}

MaxPool2 input-VJP — gradient routes only to the argmax positions.

The backward function implements:

dx[c, 2h+a, 2w+b] = dy[c, h, w] · 𝟙[(a,b) is the argmax of the window]

Conceptually, max-pool is a piecewise selection: each output is one specific input. So the Jacobian is a sparse 0/1 matrix and the VJP just routes the gradient to the chosen input.

The verified renders (the maxPoolBack node, printed by Proofs.StableHLO.pretty) emit stablehlo.select_and_scatter over the 2×2 stride-2 windows:

// select: keep the current cell unless a later one is strictly larger %sge = stablehlo.compare GE, %sa, %sb // scatter: add the routed cotangent into the chosen cell %ss = stablehlo.add %sc, %sd

The GE select keeps the first window cell attaining the max, so each window's cotangent lands in ONE cell, as PyTorch's and JAX's max-pool backwards do.

Canonical (junk-at-tie) witness. HasVJP3.correct is satisfied by the canonical pdiv3-derived backward via rfl. At argmax-tie boundaries maxPool2 is not differentiable, so pdiv3 agrees with fderiv's junk default of 0 and the canonical witness is also 0 there, where the render routes the cotangent to the first tied cell. Off ties the two agree: see maxPool2_codegen_matches_canonical below. The generic MlirCodegen walk still emits the older tile-compare-select (compare EQ of the input against the tiled window max, then select), which routes the cotangent to every tied cell. See LeanMlir/Proofs/README.md's trust-boundary section.

Named accessor for the maxPool2 input backward (the canonical witness of maxPool2HasVJP3).

Row of the 2×2 window that contains input row hi_in.

Position within the window's row (0 = top, 1 = bottom).

Row index of position a ∈ Fin 2 inside output window hi_out.

Smoothness: every 2×2 window of x attains its max at exactly one cell — a cell that dominates its window is strictly above every other cell of it. The other cells may tie with each other (a post-ReLU window [5, 0, 0, 0] qualifies); a window whose max is attained twice (an all-zero post-ReLU window) does not, and there maxPool2 has no derivative.

Input position (ci, hi_in, wi_in) attains the max of its window.

A (not necessarily unique) argmax of the 2×2 window at output position (co, ho, wo). Unique under MaxPool2Smooth.

For each output flat index k_out (decoded to (co, ho, wo)), the flat index of the argmax's input position in Vec (c * (2*h) * (2*w)). Used as the carrier of the local-linearization reindexCLM.

MaxPool2 pointwise VJP — no canonical-witness escape.

HasVJPAt3 maxPool2 x under MaxPool2Smooth x. The backward is the argmax routing directly (route dy to the argmax cell, zero elsewhere); the correct field is maxPool2_codegen_matches_canonical flipped, not rfl. Companion of reluHasVJPAt in MLP.lean: mlpHasVJPAt chains through reluHasVJPAt, and cnnHasVJPAt chains through both (this one via maxPoolFlatHasVJPAt), so neither uses the global canonical witness at a kinked operator.

Global average pool: mean over spatial dims per channel.

Flat GAP: Vec (c*h*w) → Vec c.

The channel of a flat index idx : Fin (c*h*w).

Global average pool VJP (flattened). Linear map; backward broadcasts dy ci / (h*w) to every spatial cell of channel ci.

The 2×2/s2 max pool maxPool2 on the flattened Vec form (unflatten, pool, flatten).

The VJP of maxPoolFlat at a flattened input whose 2×2 windows each have a unique max (MaxPool2Smooth): maxPool2HasVJPAt3 moved to Vec form.

The forward CNN: stem(convBnRelu) → maxpool → resblock(id) → resblockProj → gap → dense.