Documentation

LeanMlir.Proofs.Foundation.BatchMapVJPAt

`batchMap` at a POINT — the pointwise peer of `batchMapHasVJP` #

The global lift batchMapHasVJP (below) takes HasVJP f and Differentiable ℝ f; that is enough for EfficientNet-B0, whose swish is smooth everywhere and whose stem has no pooling. batchMapHasVJPAt is its pointwise peer, which ResNet-34's stem needs: the stem is batchMap N (maxPool3s2Flat c h w) ∘ cbReluStridedB, and a max-pool has no derivative at a tie, so maxPool3s2FlatHasVJPAtVec (BackwardMaps) is _at by nature.

The weakening is exactly as narrow as it looks. pdivMat_rowIndep requires Differentiable ℝ g, and its docstring explains why (a non-differentiable coordinate makes fderiv junk and breaks the per-row decomposition) — but every use of that hypothesis is at a ROW of the matrix it is stated about. So it weakens to ∀ r, DifferentiableAt ℝ g (A r), which is how Tensor.lean's pdivMat_rowIndep_perRow_at states it; pdivMat_rowIndep_at below is that lemma with the same map on every row.

The r34 stem's instance lives with r34's VJP (ResNet34FullBVJP) and in HeadLayers. It is two lines: batchMapHasVJPAt _ v (fun r => maxPool3s2FlatHasVJPAtVec (Mat.unflatten v r) (hs r)) (fun r => maxPool3s2Flat_differentiableAt_vec (Mat.unflatten v r) (hs r)), with no glue between the two, which is what says the lemma below has the right shape.

theorem Proofs.pdivMat_rowIndep_at {m n p : ℕ} (g : Vec n → Vec p) (A : Mat m n) (h_g_diff : ∀ (r : Fin m), DifferentiableAt ℝ g (A r)) (i : Fin m) (j : Fin n) (k : Fin m) (l : Fin p) :

pdivMat (fun (M : Mat m n) (r : Fin m) => g (M r)) A i j k l = if i = k then pdiv g (A i) j l else 0

Row-wise Jacobian decomposition at a point. pdivMat_rowIndep_perRow_at (Tensor.lean) with one map g for every row: global differentiability of g weakened to differentiability at each ROW of A.

theorem Proofs.batchMap_differentiableAt {N a b : ℕ} (f : Vec a → Vec b) (v : Vec (N * a)) (hf : ∀ (r : Fin N), DifferentiableAt ℝ f (Mat.unflatten v r)) :

DifferentiableAt ℝ (StableHLO.batchMap N f) v

batchMap N f is differentiable at v when f is differentiable at each of v's rows. The pointwise peer of batchMap_differentiable.

theorem Proofs.pdiv_batchMap_at {N a b : ℕ} (f : Vec a → Vec b) (v : Vec (N * a)) (hf_diff : ∀ (r : Fin N), DifferentiableAt ℝ f (Mat.unflatten v r)) (idx : Fin (N * a)) (jdx : Fin (N * b)) :

pdiv (StableHLO.batchMap N f) v idx jdx = if (finProdFinEquiv.symm idx).1 = (finProdFinEquiv.symm jdx).1 then pdiv f (Mat.unflatten v (finProdFinEquiv.symm idx).1) (finProdFinEquiv.symm idx).2 (finProdFinEquiv.symm jdx).2 else 0

batchMap's Jacobian is block-diagonal across the batch, at a point. pdivMat_rowIndep_at read through batchMap_eq_rowwiseFlat: entry (idx, jdx) vanishes unless the two indices name the same example, and is f's own entry on that example's row otherwise.

noncomputable def Proofs.batchMapHasVJPAt {N a b : ℕ} (f : Vec a → Vec b) (v : Vec (N * a)) (hf : (r : Fin N) → HasVJPAt f (Mat.unflatten v r)) (hf_diff : ∀ (r : Fin N), DifferentiableAt ℝ f (Mat.unflatten v r)) :

HasVJPAt (StableHLO.batchMap N f) v

batchMap N f's VJP at a point — the pointwise peer of batchMapHasVJP, and the lift a batch-separable op with a KINK needs. The backward reshapes to [N, ·] and runs each example's own _at backward on its own row, exactly as the global one runs f.backward row-wise.

Note: Unlike batchMapHasVJP this is built field by field rather than transported along batchMap_eq_rowwiseFlat with ▸: an Eq.mpr blocks .backward from reducing, which a whole-net certified-backward tie later needs.

Equations

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

Instances For

theorem Proofs.batchMapAux_eq_batchMapHasVJPAt {N a b : ℕ} (f : Vec a → Vec b) (g : Vec a → Vec b → Vec a) (v : Vec (N * a)) (hf : (r : Fin N) → HasVJPAt f (Mat.unflatten v r)) (hf_diff : ∀ (r : Fin N), DifferentiableAt ℝ f (Mat.unflatten v r)) (hg : ∀ (r : Fin N), g (Mat.unflatten v r) = (hf r).backward) :

StableHLO.batchMapAux N g v = (batchMapHasVJPAt f v hf hf_diff).backward

A batched backward tie from the per-example one. If g is each row's certified backward, batchMapAux N g v IS the lifted witness's backward: both read example r's row of v and dy, so the two agree index by index once the rows do.

theorem Proofs.batchMap_eq_batchMapHasVJPAt {N a b : ℕ} (f : Vec a → Vec b) (g : Vec b → Vec a) (v : Vec (N * a)) (hf : (r : Fin N) → HasVJPAt f (Mat.unflatten v r)) (hf_diff : ∀ (r : Fin N), DifferentiableAt ℝ f (Mat.unflatten v r)) (hg : ∀ (r : Fin N), g = (hf r).backward) :

StableHLO.batchMap N g = (batchMapHasVJPAt f v hf hf_diff).backward

The linear-backward form of batchMapAux_eq_batchMapHasVJPAt: a backward that ignores the saved input lifts by batchMap.

theorem Proofs.batchMap_comp (B : ℕ) {a b c : ℕ} (f : Vec a → Vec b) (g : Vec b → Vec c) :

StableHLO.batchMap B (g ∘ f) = StableHLO.batchMap B g ∘ StableHLO.batchMap B f

batchMap B (g ∘ f) = batchMap B g ∘ batchMap B f. Both sides read example p.1's slice of the input and run g ∘ f on it; peeling the inner lift off at one example is batchSlice_batchMap. The two spellings are NOT rfl — they agree only up to finProdFinEquiv.symm_apply_apply — which is why every batched whole-net chain saves its activations stage by stage (vitSavedBodyB, cnxSavedB1 … cnxSavedB10) and its shape check goes through this lemma. Shared by the ViT and ConvNeXt batched ties.

theorem Proofs.batchMap_eq_rowwiseFlat {N a b : ℕ} (f : Vec a → Vec b) :

StableHLO.batchMap N f = fun (v : Vec (N * a)) => Mat.flatten ((fun (A : Mat N a) (r : Fin N) => f (A r)) (Mat.unflatten v))

batchMap N f is the flattened row-wise application of f. Reading the output at flat index idx (decoding to example m, coord c) gives f (row m of the input) c on both sides — the Mat.flatten/unflatten row-major convention is exactly batchMap's finProdFinEquiv split.

theorem Proofs.batchMap_differentiable {N a b : ℕ} (f : Vec a → Vec b) (hf : Differentiable ℝ f) :

Differentiable ℝ (StableHLO.batchMap N f)

batchMap N f is differentiable when f is — it is f applied independently per example.

noncomputable def Proofs.batchMapHasVJP {N a b : ℕ} (f : Vec a → Vec b) (hf : HasVJP f) (hf_diff : Differentiable ℝ f) :

HasVJP (StableHLO.batchMap N f)

batchMap N f VJP — block-diagonal (the genuinely-new lemma). A batch-separable op's VJP applies f's proven VJP independently per example. The backward, like the forward, reshapes to [N, ·] and runs f.backward row-wise. Reuses rowwiseHasVJPMat + HasVJPMat.toHasVJP. This is seBlockFullHasVJP / the conv-depthwise-dense VJPs "lifted by batchMap" to the whole batch.

Equations

Proofs.batchMapHasVJP f hf hf_diff = Proofs.HasVJP.congr ⋯ (Proofs.rowwiseHasVJPMat hf hf_diff).toHasVJP

Instances For

noncomputable def Proofs.reindexHasVJP {a b : ℕ} (σ : Fin b → Fin a) :

HasVJP ⇑(reindexCLM σ)

Reindex VJP at reindexCLM — reindexVJP, with the backward spelled as the masked sum the IR's scatter ops denote (reindexVJP_backward), so the batched graph ties close by rfl.

Equations

Proofs.reindexHasVJP σ = { backward := fun (_v : Proofs.Vec a) (dy : Proofs.Vec b) (i : Fin a) => ∑ k : Fin b, if i = σ k then dy k else 0, correct := ⋯ }

Instances For

theorem Proofs.bnBatchLA_eq_comp (N oc h w : ℕ) (ε : ℝ) (γ β : Vec oc) :

StableHLO.bnBatchLA N oc h w ε γ β = ⇑(reindexCLM (Fin.cast ⋯)) ∘ bnBatchTensor4 N oc h w ε γ β ∘ ⇑(reindexCLM (Fin.cast ⋯))

bnBatchLA is the proven bnBatchTensor4, conjugated by the mul_assoc reindex. Both reindex maps are reindexCLM (Fin.cast …); the middle is the genuinely batch-coupled true batch-norm.

theorem Proofs.bnBatchLA_differentiable (N oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

Differentiable ℝ (StableHLO.bnBatchLA N oc h w ε γ β)

theorem Proofs.bnBatchLA_continuous (N oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

Continuous (StableHLO.bnBatchLA N oc h w ε γ β)

noncomputable def Proofs.bnBatchLAHasVJP (N oc h w : ℕ) (ε : ℝ) (hε : 0 < ε) (γ β : Vec oc) :

HasVJP (StableHLO.bnBatchLA N oc h w ε γ β)

True batch-norm VJP at the network's flat index. bnBatchLA's backward is the proven bnBatchTensor4 VJP (batch-coupled — NOT a batchMap), conjugated by the reindex isos.

Equations

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

Instances For