A depthwise kernel: (c, kH, kW) — one filter per channel, no in_channels axis. (Equivalently, a (c, 1, kH, kW) kernel where the singleton dim has been squeezed out.)

In the MLIR backend this is represented as a regular (c, 1, kH, kW) kernel with feature_group_count = c, telling StableHLO to apply each kernel only to its own input channel.

Depthwise conv2d forward (SAME padding, stride 1).

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

Compare to regular conv2d (CNN.lean conv2d): regular: y[o,h,w] = Σ_{c, kh, kw} x[c,...] · W[o,c,kh,kw] + b[o] depthwise: same minus the Σ_c (no cross-channel mixing).

Output has the same number of channels as the input (c, not oc).

MLIR (MlirCodegen.lean emitDepthwiseConvBn): uses feature_group_count = c to tell StableHLO that each kernel applies only within its own channel group.

Closed-form input gradient for depthwise conv2d — direct formula, written as a sum over output positions (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). No Σ co like regular conv2d — input channel ci reads only from kernel-channel ci and gradient-channel ci, because depthwise has no cross-channel mixing.

Equivalent (under the (ho, wo) ↔ (kh, kw) partial bijection) to the MLIR-aligned reversed-kernel formula dx[c, h, w] = Σ_{kh, kw} W[c, kH−1−kh, kW−1−kw] · dy[c, h+kh−p, w+kw−p].

Depthwise conv input-VJP — proved from foundation rules.

The function v ↦ flatten (depthwiseConv2d W b (unflatten v)) is affine in v: a constant b ohw_o(idx_out) plus a double sum over (kh, kw) of W ohw_o kh kw * (if pad-cond then v(reindex) else 0). Mirrors conv2dHasVJP3 but with one fewer sum level (no Σ c) and the channel for the v-read is the same as ohw_o (forced by structure: input-channel = output-channel in depthwise).

The closing collapse first folds Σ co → co=ci (since for co ≠ ci, the indicator idx_in = finProdFinEquiv (..., co, ...) is false by channel-mismatch on the first projection), then proceeds per-(ho, wo) with a 2-conjunct h_indicator (just kh+ho = hi+pH and kw+wo = wi+pW; no c = ci since c isn't summed).

The backward function (accessed as (depthwiseHasVJP3 W b).backward, or via the depthwiseConv2dInputGrad abbrev below) implements depthwiseConv2dInputGradFormula W dy ci hi wi. Equivalent to the MLIR-aligned reversed-kernel formula dx[c, h, w] = Σ_{kh, kw} W[c, kH−1−kh, kW−1−kw] · dy[c, h+kh−p, w+kw−p].

Named accessor for the depthwise input backward — aligns with MLIR codegen (per-channel stablehlo.convolution in the backward pass).

Flat depthwise conv — depthwiseConv2d bridged into flattened Vec → Vec space: flatten ∘ depthwiseConv2d W b ∘ unflatten. Channels and spatial dims are preserved (c h w → c h w), so this is Vec (c*h*w) → Vec (c*h*w). Mirror of flatConv; the form the MobileNet/EfficientNet/ConvNeXt VJP composition uses (flat Vec space).

Flat depthwise conv input-VJP. depthwiseFlat W b is defeq to the generic bridge's fun v => flatten (depthwiseConv2d W b (unflatten v)), so HasVJP3.toHasVJP applied to depthwiseHasVJP3 lands the witness directly. Mirror of the regular-conv flat VJP.

Stride-2 SAME depthwise conv, flattened: Vec (c·2h·2w) → Vec (c·h·w). Defined as decimateFlat ∘ depthwiseFlat (the stride-1 SAME depthwise on the 2h×2w grid, then keep even positions) — exactly the strided-conv recipe (flatConvStride2, StridedConv.lean) with the depthwise kernel. This is how MobileNetV2 downsamples (stride-2 depthwise inside an inverted-residual block); channels are unchanged (c → c), spatial halves.

Stride-2 depthwise input-VJP — by the chain rule (vjpComp) on decimateFlat ∘ depthwiseFlat, reusing the proven stride-1 depthwise input-VJP (depthwiseFlatHasVJP) and the decimation VJP. The backward is depthwise.back (decimate.back dy) — i.e. zero-upsample the cotangent then run the reversed-kernel stride-1 depthwise (StableHLO: stablehlo.pad interior=1 then feature_group_count = c reversed-kernel conv), exactly the convStridedBack shape with the per-channel grouping.

Depthwise weight-VJP — proved from foundation rules.

DepthwiseKernel c kH kW is definitionally Tensor3 c kH kW, so HasVJP3 applies directly. The function W ↦ depthwiseConv2d W b x is affine in W: at output (co, ho, wo) it's b co + Σ_{kh, kw} W co kh kw * x_pad_term(co, kh, kw, ho, wo). Same recipe as conv2dWeightGradHasVJP but with two inner dims (kh, kw) instead of three (c, kh, kw) — depthwise has no cross-channel sum, so the "channel match" condition co = ci is a single equality rather than a packed comparison.

Named accessor for the depthwise weight backward.

Depthwise bias-VJP — proved from foundation rules. Same shape as conv2dBiasGradHasVJP, just simpler: depthwise has no Σ over input channels (input channel = output channel). The function b ↦ flatten(depthwiseConv2d W b x) decomposes as (channel-reindex from b) + (W,x term constant in b), exactly like conv2d's case.

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

Depthwise bias gradient — closed-form formula.

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

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

Identical to regular conv's bias gradient — the bias is per-channel in both cases, and it adds the same value to every spatial cell of its channel. The reduction is the same.

Stride-2 depthwise weight-VJP. fun v => depthwiseStride2Flat (unflatten v) b x = decimate ∘ (depthwise-weight-in-v); by vjpComp of the proven stride-1 depthwiseWeightGradHasVJP3 (flattened via HasVJP3.toHasVJP) with decimateFlatHasVJP. The depthwise peer of flatConvStride2WeightGradHasVJP.

Stride-2 depthwise bias-VJP. fun b => depthwiseStride2Flat W b x = decimate ∘ (depthwise-bias-in-b); by vjpComp of the proven stride-1 depthwiseBiasGradHasVJP with decimateFlatHasVJP.

Stride-2 XLA-SAME depthwise conv, flattened: Vec (c·2h·2w) → Vec (c·h·w). decimateOddFlat ∘ depthwiseFlat — the asymmetric-pad peer of depthwiseStride2Flat.

Stride-2 XLA-SAME depthwise input-VJP. vjpComp on decimateOddFlat ∘ depthwiseFlat. The backward zero-upsamples the cotangent onto the odd positions, then runs the reversed-kernel grouped conv — so the forward's asymmetry is placed by the backward too. A symmetric backward against this forward is a silent wrong-gradient.

Stride-2 XLA-SAME depthwise weight-VJP. The kernel-side peer, by vjpComp of the proven stride-1 depthwiseWeightGradHasVJP3 with the odd-decimation VJP.

Stride-2 XLA-SAME depthwise bias-VJP.

Stride-1 depthwise weight SGD step: flatten W − lr·flatten(dwconv_weight_grad(b,x)·dy).

Stride-1 depthwise bias SGD step: b − lr·(dwconv_bias_grad(W,x)·dy).

Stride-2 depthwise weight SGD step: flatten W − lr·(dwconvStride2_weight_grad(b,x)·dy).

Stride-2 depthwise bias SGD step: b − lr·(dwconvStride2_bias_grad(W,x)·dy).

Stride-2 XLA-SAME depthwise weight SGD step — depthwiseStridedWeightSgdDen's peer at the odd decimation phase (depthwiseStride2XlaWeightGradHasVJP). Same non-reducing wrapper, for the same den-match-size reason.

Stride-2 XLA-SAME depthwise bias SGD step.