Documentation

LeanMlir.Proofs.Foundation.BatchMapVJPAt

batchMap at a POINT — the pointwise peer of batchMap_has_vjp #

EfficientNetChainClose.lean lifts a batch-separable op's VJP to the whole batch in the GLOBAL form: batchMap_has_vjp takes HasVJP f and Differentiable ℝ f. EfficientNet never needed anything weaker — swish is smooth everywhere and B0's stem has no pooling — so the pointwise peer was never written.

ResNet-34 needs it. Its stem is batchMap N (maxPool3s2Flat c h w) ∘ cbReluStridedB, and a max-pool has no derivative at a tie: maxPool3s2Flat_has_vjp_at_vec is _at by nature. Without the lift below, r34's whole-net VJP at batch BN cannot be assembled — the one thing standing between ResNet34FullB.lean and T1.

⭐ 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, not here — maxPool3s2Flat_has_vjp_at_vec is in the Float tier and this is a Foundation file. It is two lines there: batchMap_has_vjp_at _ v (fun r => maxPool3s2Flat_has_vjp_at_vec (Mat.unflatten v r) (hs r)) (fun r => maxPool3s2Flat_differentiableAt_vec (Mat.unflatten v r) (hs r) hc hh hw), with no glue between the two, which is what says the lemma below has the right shape.

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

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

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

source
noncomputable def Proofs.batchMap_has_vjp_at {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 batchMap_has_vjp, 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.

⚠ Unlike batchMap_has_vjp 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
    source
    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.