Depthwise Convolution

The structural simplification at the heart of MobileNet, EfficientNet, and every "mobile" CNN since ~2017. Standard conv2d does cross-channel mixing (every output channel sees every input channel) plus spatial filtering. Depthwise conv drops the cross-channel mixing: each input channel gets its own 2D filter and produces its own output channel.

Math-wise it's "regular conv with a constraint." Practical-wise it's ~10× cheaper because you avoid the O(ic · oc) cross-channel sum.

Architecturally, depthwise is always paired with a 1×1 "pointwise" conv that does the cross-channel mixing separately. Together they form the depthwise-separable convolution (Xception, MobileNet) — the same expressive power as a regular conv, factored into two cheaper steps.

What this file proves

The depthwise conv is structurally a special case of regular conv:

• Regular conv kernel: (oc, ic, kH, kW) — full mixing.
• Depthwise kernel: (c, 1, kH, kW) — diagonal in the channel pair.

So we don't re-derive the VJPs from scratch. We state them as the "channel-restricted" versions of conv2d_input_grad / conv2d_weight_grad from CNN.lean. The transpose trick still works, the reversed-kernel trick still works — they just operate per-channel.

@[reducible, inline]

abbrev Proofs.DepthwiseKernel (c kH kW : ℕ) : Type

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.

Equations

• Proofs.DepthwiseKernel c kH kW = (Fin c → Fin kH → Fin kW → ℝ)

noncomputable def Proofs.depthwiseConv2d {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (x : Tensor3 c h w) : Tensor3 c h w

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.

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.depthwiseConv2d_input_grad_formula {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (dy : Tensor3 c h w) : Tensor3 c h w

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].

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.depthwise_has_vjp3 {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) : HasVJP3 (depthwiseConv2d W b)

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 conv2d_has_vjp3 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 (depthwise_has_vjp3 W b).backward, or via the depthwiseConv2d_input_grad abbrev below) implements depthwiseConv2d_input_grad_formula 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].

Equations

• Proofs.depthwise_has_vjp3 W b = { backward := fun (_x dy : Proofs.Tensor3 c h w) => Proofs.depthwiseConv2d_input_grad_formula W dy, correct := ⋯ }

@[reducible, inline]

noncomputable abbrev Proofs.depthwiseConv2d_input_grad {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (x dy : Tensor3 c h w) : Tensor3 c h w

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

Equations

• Proofs.depthwiseConv2d_input_grad W b x dy = (Proofs.depthwise_has_vjp3 W b).backward x dy

theorem Proofs.depthwise_differentiable {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) : Differentiable ℝ (depthwiseConv2d W b)

depthwiseConv2d is differentiable everywhere. Mirror of conv2d_differentiable: depthwiseConv2d W b x ch hi wi is the affine map b ch + ∑_{kh,kw} W ch kh kw · (pad-eval x) — a constant bias plus a finite ℝ-linear combination of input coordinates (the dependent if-pad- eval being a projection or the constant 0). differentiable_pi reduces to per-coordinate differentiability; DifferentiableAt.fun_sum lifts the double sum (no Σ c — depthwise reads only its own channel).

noncomputable def Proofs.depthwiseFlat {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) : Vec (c * h * w) → Vec (c * h * w)

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).

Equations

• Proofs.depthwiseFlat W b v = (Proofs.depthwiseConv2d W b (Proofs.Tensor3.unflatten v)).flatten

theorem Proofs.depthwiseFlat_differentiable {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) : Differentiable ℝ (depthwiseFlat W b)

depthwiseFlat is differentiable everywhere. Composition of the three differentiable maps unflatten, depthwiseConv2d, flatten. Mirror of flatConv_differentiable.

noncomputable def Proofs.depthwiseFlat_has_vjp {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) : HasVJP (depthwiseFlat W b)

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_to_hasVJP applied to depthwise_has_vjp3 lands the witness directly. Mirror of the regular-conv flat VJP.

Equations

• Proofs.depthwiseFlat_has_vjp W b = Proofs.hasVJP3_to_hasVJP (Proofs.depthwise_has_vjp3 W b)

noncomputable def Proofs.depthwiseStride2Flat {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) : Vec (c * (2 * h) * (2 * w)) → Vec (c * h * w)

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.

Equations

• Proofs.depthwiseStride2Flat W b = Proofs.decimateFlat c h w ∘ Proofs.depthwiseFlat W b

theorem Proofs.depthwiseStride2Flat_differentiable {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) : Differentiable ℝ (depthwiseStride2Flat W b)

noncomputable def Proofs.depthwiseStride2Flat_has_vjp {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) : HasVJP (depthwiseStride2Flat W b)

Stride-2 depthwise input-VJP — by the chain rule (vjp_comp) on decimateFlat ∘ depthwiseFlat, reusing the proven stride-1 depthwise input-VJP (depthwiseFlat_has_vjp) 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.

Equations

• Proofs.depthwiseStride2Flat_has_vjp W b = Proofs.vjp_comp (Proofs.depthwiseFlat W b) (Proofs.decimateFlat c h w) ⋯ ⋯ (Proofs.depthwiseFlat_has_vjp W b) (Proofs.decimateFlat_has_vjp c h w)

theorem Proofs.depthwiseStride2Flat_has_vjp_correct {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (x : Vec (c * (2 * h) * (2 * w))) (dy : Vec (c * h * w)) (i : Fin (c * (2 * h) * (2 * w))) : (depthwiseStride2Flat_has_vjp W b).backward x dy i = ∑ j : Fin (c * h * w), pdiv (depthwiseStride2Flat W b) x i j * dy j

Stride-2 depthwise input-VJP correctness (the ℝ-carrying audit headline): the backward equals the pdiv-contracted Jacobian of depthwiseStride2Flat.

Depthwise weight gradient (Phase 7 — proved from foundation rules)

Per-channel transpose trick:

`dW[c, kh, kw] = Σ_{h, w} x[c, h+kh−p, w+kw−p] · dy[c, h, w]`

Compare to the regular conv weight gradient:

• Regular conv: produces (oc, ic, kH, kW) — every (oc, ic) pair.
• Depthwise: produces (c, kH, kW) — only the diagonal (c, c) pairs survive (the rest are zero by construction).

The transpose trick works the same way: view x and dy with channel and batch axes swapped, do a standard conv, the spatial dims of dy become the kernel dims. The only difference is that feature_group_count is set so the conv stays per-channel.

MLIR (the depthwise variant of the transpose trick is in emitDepthwiseConvBnBackward around line 1855):

"For depthwise: dW[c,1,kH,kW] = sum_b input[b,c,:,:] conv grad[b,c,:,:]"

Framework. Unlike the regular-conv weight gradient (which needs Kernel4.flatten because the kernel is 4D), the depthwise kernel DepthwiseKernel c kH kW is 3D — same shape as Tensor3 c kH kW, and in fact definitionally equal. So we can reuse the existing HasVJP3 framework directly, parameterized over W instead of x.

noncomputable def Proofs.depthwise_weight_grad_has_vjp3 {c h w kH kW : ℕ} (b : Vec c) (x : Tensor3 c h w) : HasVJP3 fun (W : DepthwiseKernel c kH kW) => depthwiseConv2d W b x

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 conv2d_weight_grad_has_vjp 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.

Equations

• One or more equations did not get rendered due to their size.

@[reducible, inline]

noncomputable abbrev Proofs.depthwiseConv2d_weight_grad {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (x dy : Tensor3 c h w) : DepthwiseKernel c kH kW

Named accessor for the depthwise weight backward.

Equations

• Proofs.depthwiseConv2d_weight_grad W b x dy = (Proofs.depthwise_weight_grad_has_vjp3 b x).backward W dy

noncomputable def Proofs.depthwise_bias_grad_has_vjp {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (x : Tensor3 c h w) : HasVJP fun (b : Vec c) => (depthwiseConv2d W b x).flatten

Depthwise bias-VJP — proved from foundation rules. Same shape as conv2d_bias_grad_has_vjp, 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.

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.depthwiseConv2d_bias_grad {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (x dy : Tensor3 c h w) : Vec c

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

Equations

• Proofs.depthwiseConv2d_bias_grad W b x dy = (Proofs.depthwise_bias_grad_has_vjp W x).backward b dy.flatten

noncomputable def Proofs.depthwiseConv2d_bias_grad_formula {c h w : ℕ} (dy : Tensor3 c h w) : Vec c

Depthwise bias gradient — closed-form formula (documented, numerically verified, expected to equal depthwiseConv2d_bias_grad up to fp precision).

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

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.

Equations

• Proofs.depthwiseConv2d_bias_grad_formula dy cc = ∑ y : Fin h, ∑ x : Fin w, dy cc y x

Depthwise = constrained regular conv

Conceptually, depthwise conv is regular conv with a sparsity pattern on the kernel: W_regular[o, c, kh, kw] is zero unless o = c. Equivalently, W_regular is block-diagonal in the (o, c) channel pair.

Two consequences for the VJPs:

1. Forward: the cross-channel sum Σ_c collapses to a single term (the diagonal one), giving the per-channel formula above.

2. Backward: every formula involving Σ_o or Σ_c over channel indices collapses similarly. The transpose trick still produces a (ic, oc, kH, kW) tensor in principle, but only the diagonal slice is nonzero, and the implementation just stores the diagonal.

So you don't have to derive depthwise VJPs from scratch — you derive them by specializing the regular conv VJPs to the sparsity pattern. This is a great example of how a constraint on the forward propagates mechanically to a constraint on the backward.

Cost

The forward op cost goes from O(B · oc · ic · H · W · kH · kW) for regular conv to O(B · c · H · W · kH · kW) for depthwise — saves a factor of oc (typically 32–512). The backward cost reduces by the same factor. This is why mobile architectures pair depthwise with a cheap 1×1 pointwise conv for cross-channel mixing — together they have the same expressive power as a regular conv at a fraction of the FLOPs.

Where it's used

• MobileNet v1/v2/v3 (MainMobilenet.lean, MainMobilenetV2.lean, MainMobilenetV3.lean) — depthwise everywhere.
• EfficientNet (MainEfficientNet.lean) — depthwise inside MBConv blocks.
• MBConv (MainEfficientNet.lean, MainEfficientNetV2.lean) — the block that pairs an expand 1×1 → depthwise k×k → project 1×1, with optional Squeeze-and-Excitation. See SE.lean for the SE part.

Summary of derivations in this file

None. The forward depthwiseConv2d is a concrete definition (not a black-box), and all three VJPs are theorems proved from the foundation rules in Tensor.lean:

• depthwise_has_vjp3 — input-path VJP. Phase 2 (Apr 2026): proved from pdiv_add / pdiv_const / pdiv_finset_sum / pdiv_const_mul_pi_pad_eval. Mirrors conv2d_has_vjp3 with one fewer sum level (no Σ c) and a prepended Σ co collapse.
• depthwise_weight_grad_has_vjp3 — weight-path VJP, bundled as HasVJP3 directly (no flattening needed; see framework note above). Gradient-checked numerically.
• depthwise_bias_grad_has_vjp — bias-path VJP, bundled HasVJP on the flattened output. Same pattern as conv2d's bias VJP.

Pure-Mathlib closure verified via #print axioms (only propext, Classical.choice, Quot.sound).

Derived helpers (not axioms):

• depthwiseConv2d_input_grad, depthwiseConv2d_weight_grad, depthwiseConv2d_bias_grad — named accessors, .backward of the corresponding VJP.
• depthwiseConv2d_input_grad_formula — the concrete sum-over-output- positions closed-form, used as the backward of depthwise_has_vjp3.
• depthwiseConv2d_bias_grad_formula — the concrete sum-over-spatial closed-form (numerically verified to equal the bias-VJP's backward).

theorem Proofs.depthwise_has_vjp3_correct {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (x dy : Tensor3 c h w) (ci : Fin c) (hi : Fin h) (wi : Fin w) : (depthwise_has_vjp3 W b).backward x dy ci hi wi = ∑ co : Fin c, ∑ ho : Fin h, ∑ wo : Fin w, pdiv3 (depthwiseConv2d W b) x ci hi wi co ho wo * dy co ho wo

Public correctness theorem for depthwise_has_vjp3: the proved input-VJP's backward equals the pdiv3-contracted Jacobian.

theorem Proofs.depthwise_weight_differentiable {c h w kH kW : ℕ} (b : Vec c) (x : Tensor3 c h w) : Differentiable ℝ fun (v : Vec (c * kH * kW)) => (depthwiseConv2d (Tensor3.unflatten v) b x).flatten

depthwiseConv2d (as a function of its kernel) is differentiable — affine in W. The depthwise peer of conv2d_weight_differentiable; the vjp_comp hypothesis for the strided weight-grad.

theorem Proofs.depthwise_bias_differentiable {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (x : Tensor3 c h w) : Differentiable ℝ fun (b : Vec c) => (depthwiseConv2d W b x).flatten

depthwiseConv2d (as a function of its bias) is differentiable — affine in b. The vjp_comp hypothesis for the strided depthwise bias-grad.

noncomputable def Proofs.depthwiseStride2_weight_grad_has_vjp {c h w kH kW : ℕ} (b : Vec c) (x : Vec (c * (2 * h) * (2 * w))) : HasVJP fun (v : Vec (c * kH * kW)) => depthwiseStride2Flat (Tensor3.unflatten v) b x

Stride-2 depthwise weight-VJP. fun v => depthwiseStride2Flat (unflatten v) b x = decimate ∘ (depthwise-weight-in-v); by vjp_comp of the proven stride-1 depthwise_weight_grad_has_vjp3 (flattened via hasVJP3_to_hasVJP) with decimateFlat_has_vjp. The depthwise peer of flatConvStride2_weight_grad_has_vjp.

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.depthwiseStride2_bias_grad_has_vjp {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (x : Vec (c * (2 * h) * (2 * w))) : HasVJP fun (b : Vec c) => depthwiseStride2Flat W b x

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

Equations

• One or more equations did not get rendered due to their size.

The depthwise peer of flatConvStride2Xla (Foundation/StridedConv.lean), and it exists for the same reason: jax/Jax/Codegen.lean:679's depthwise_conv defaults to padding='SAME', so MobileNetV2's four strided depthwises — and EfficientNet's — pad asymmetrically, while depthwiseStride2Flat above pads symmetrically. Both give the same output size, so only a forward tie can see it; planning/archive/mnv4_verified.md §3d measured MNv2's five sites at 2.9e-1 of a ~1.05 logit range in its trainer's BN world.

⭐ Identical structure to the regular-conv case, so identical cost: the asymmetry is a phase shift in the decimation, decimateOddFlat instead of decimateFlat. depthwiseFlat and all of its VJPs are reused verbatim, and decimateOddFlat_has_vjp is already proven, so nothing here is a new obligation.

⚠ Even inputs only — which the type enforces (c*(2*h)*(2*w)) and which is every strided depthwise in mnv2/mnv4/enet (112, 56, 28, 14). At an odd input XLA SAME is symmetric and depthwiseStride2Flat is already correct.

noncomputable def Proofs.depthwiseStride2FlatXla {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) : Vec (c * (2 * h) * (2 * w)) → Vec (c * h * w)

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

Equations

• Proofs.depthwiseStride2FlatXla W b = Proofs.decimateOddFlat c h w ∘ Proofs.depthwiseFlat W b

theorem Proofs.depthwiseStride2FlatXla_differentiable {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) : Differentiable ℝ (depthwiseStride2FlatXla W b)

noncomputable def Proofs.depthwiseStride2FlatXla_has_vjp {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) : HasVJP (depthwiseStride2FlatXla W b)

Stride-2 XLA-SAME depthwise input-VJP. vjp_comp 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.

Equations

• One or more equations did not get rendered due to their size.

theorem Proofs.depthwiseStride2FlatXla_has_vjp_correct {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (b : Vec c) (x : Vec (c * (2 * h) * (2 * w))) (dy : Vec (c * h * w)) (i : Fin (c * (2 * h) * (2 * w))) : (depthwiseStride2FlatXla_has_vjp W b).backward x dy i = ∑ j : Fin (c * h * w), pdiv (depthwiseStride2FlatXla W b) x i j * dy j

Stride-2 XLA-SAME depthwise input-VJP correctness (the ℝ-carrying audit headline).

noncomputable def Proofs.depthwiseStride2Xla_weight_grad_has_vjp {c h w kH kW : ℕ} (b : Vec c) (x : Vec (c * (2 * h) * (2 * w))) : HasVJP fun (v : Vec (c * kH * kW)) => depthwiseStride2FlatXla (Tensor3.unflatten v) b x

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

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.depthwiseStride2Xla_bias_grad_has_vjp {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (x : Vec (c * (2 * h) * (2 * w))) : HasVJP fun (b : Vec c) => depthwiseStride2FlatXla W b x

Stride-2 XLA-SAME depthwise bias-VJP.

Equations

• One or more equations did not get rendered due to their size.

noncomputable def Proofs.depthwiseWeightSgdDen {c h w kH kW : ℕ} (b : Vec c) (x : Tensor3 c h w) (W : DepthwiseKernel c kH kW) (lr : ℝ) (dy : Vec (c * h * w)) : Vec (c * kH * kW)

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

Equations

• Proofs.depthwiseWeightSgdDen b x W lr dy idx = Proofs.Tensor3.flatten W idx - lr * ((Proofs.depthwise_weight_grad_has_vjp3 b x).backward W (Proofs.Tensor3.unflatten dy)).flatten idx

noncomputable def Proofs.depthwiseBiasSgdDen {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (x : Tensor3 c h w) (b : Vec c) (lr : ℝ) (dy : Vec (c * h * w)) : Vec c

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

Equations

• Proofs.depthwiseBiasSgdDen W x b lr dy o = b o - lr * (Proofs.depthwise_bias_grad_has_vjp W x).backward b dy o

noncomputable def Proofs.depthwiseStridedWeightSgdDen {c h w kH kW : ℕ} (b : Vec c) (x : Vec (c * (2 * h) * (2 * w))) (W : DepthwiseKernel c kH kW) (lr : ℝ) (dy : Vec (c * h * w)) : Vec (c * kH * kW)

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

Equations

• Proofs.depthwiseStridedWeightSgdDen b x W lr dy idx = Proofs.Tensor3.flatten W idx - lr * (Proofs.depthwiseStride2_weight_grad_has_vjp b x).backward (Proofs.Tensor3.flatten W) dy idx

noncomputable def Proofs.depthwiseStridedBiasSgdDen {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (x : Vec (c * (2 * h) * (2 * w))) (b : Vec c) (lr : ℝ) (dy : Vec (c * h * w)) : Vec c

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

Equations

• Proofs.depthwiseStridedBiasSgdDen W x b lr dy o = b o - lr * (Proofs.depthwiseStride2_bias_grad_has_vjp W x).backward b dy o

noncomputable def Proofs.depthwiseStridedXlaWeightSgdDen {c h w kH kW : ℕ} (b : Vec c) (x : Vec (c * (2 * h) * (2 * w))) (W : DepthwiseKernel c kH kW) (lr : ℝ) (dy : Vec (c * h * w)) : Vec (c * kH * kW)

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

Equations

• Proofs.depthwiseStridedXlaWeightSgdDen b x W lr dy idx = Proofs.Tensor3.flatten W idx - lr * (Proofs.depthwiseStride2Xla_weight_grad_has_vjp b x).backward (Proofs.Tensor3.flatten W) dy idx

noncomputable def Proofs.depthwiseStridedXlaBiasSgdDen {c h w kH kW : ℕ} (W : DepthwiseKernel c kH kW) (x : Vec (c * (2 * h) * (2 * w))) (b : Vec c) (lr : ℝ) (dy : Vec (c * h * w)) : Vec c

Stride-2 XLA-SAME depthwise bias SGD step.

Equations

• Proofs.depthwiseStridedXlaBiasSgdDen W x b lr dy o = b o - lr * (Proofs.depthwiseStride2Xla_bias_grad_has_vjp W x).backward b dy o