Tensor Algebra for VJP Proofs

Vectors, matrices, and operations over ℝ, using Mathlib's Finset.sum.

Partial derivatives (pdiv) and their composition rules (chain rule, linearity, product rule) are now defined and proved from Mathlib's Fréchet derivative fderiv. The post-foundation-flip definition is

pdiv f x i j := fderiv ℝ f x (basisVec i) j

and every former axiom (pdiv_id, pdiv_const, pdiv_reindex, pdiv_add, pdiv_comp, pdiv_mul) is now a theorem proved against Mathlib's API. The bilinear rules carry Differentiable hypotheses that propagate through every downstream chapter.

The post-flip path: every claim downstream of this file is either a definition Lean unfolds or a theorem typechecked against Mathlib — no project axioms remain. #print axioms vit_full_has_vjp lists only Lean core (propext, Classical.choice, Quot.sound).

@[reducible, inline]

abbrev Proofs.Vec (n : ℕ) : Type

Equations

• Proofs.Vec n = (Fin n → ℝ)

@[reducible, inline]

abbrev Proofs.Mat (m n : ℕ) : Type

Equations

• Proofs.Mat m n = (Fin m → Fin n → ℝ)

noncomputable def Proofs.Mat.mulVec {m n : ℕ} (A : Mat m n) (v : Vec n) : Vec m

Equations

• A.mulVec v i = ∑ j : Fin n, A i j * v j

def Proofs.Mat.outer {m n : ℕ} (u : Vec m) (v : Vec n) : Mat m n

Equations

• Proofs.Mat.outer u v i j = u i * v j

noncomputable def Proofs.Mat.mul {m n p : ℕ} (A : Mat m n) (B : Mat n p) : Mat m p

Equations

• A.mul B i k = ∑ j : Fin n, A i j * B j k

def Proofs.Mat.transpose {m n : ℕ} (A : Mat m n) : Mat n m

Matrix transpose: swap rows and columns.

Equations

• A.transpose j i = A i j

@[reducible]

def Proofs.basisVec {m : ℕ} (i : Fin m) : Vec m

Standard basis vector eᵢ in Vec m: 1 at index i, 0 elsewhere. Avoids Pi.single's dependent-type elaboration friction in contexts where the codomain family isn't immediately apparent.

Equations

• Proofs.basisVec i k = if k = i then 1 else 0

@[simp]

theorem Proofs.basisVec_apply {m : ℕ} (i j : Fin m) : basisVec i j = if j = i then 1 else 0

noncomputable def Proofs.reindexCLM {a b : ℕ} (σ : Fin b → Fin a) : Vec a →L[ℝ] Vec b

The reindex map y ↦ (k ↦ y (σ k)) packaged as a continuous linear map. Used to discharge pdiv_reindex and to provide DifferentiableAt evidence for reindex-shaped subexpressions.

Equations

• Proofs.reindexCLM σ = { toFun := fun (y : Proofs.Vec a) (k : Fin b) => y (σ k), map_add' := ⋯, map_smul' := ⋯, cont := ⋯ }

@[simp]

theorem Proofs.reindexCLM_apply {a b : ℕ} (σ : Fin b → Fin a) (y : Vec a) : (reindexCLM σ) y = fun (k : Fin b) => y (σ k)

noncomputable def Proofs.pdiv {m n : ℕ} (f : Vec m → Vec n) (x : Vec m) (i : Fin m) (j : Fin n) : ℝ

Partial derivative. The (i, j) entry of the Jacobian of f : Vec m → Vec n at x.

Equations

• Proofs.pdiv f x i j = (fderiv ℝ f x) (Proofs.basisVec i) j

theorem Proofs.pdiv_id {n : ℕ} (x : Vec n) (i j : Fin n) : pdiv (fun (y : Vec n) => y) x i j = if i = j then 1 else 0

Identity Jacobian — δᵢⱼ.

theorem Proofs.pdiv_const {m n : ℕ} (c : Vec n) (x : Vec m) (i : Fin m) (j : Fin n) : pdiv (fun (x : Vec m) => c) x i j = 0

Constant function Jacobian — zero.

theorem Proofs.pdiv_reindex {a b : ℕ} (σ : Fin b → Fin a) (x : Vec a) (i : Fin a) (j : Fin b) : pdiv (fun (y : Vec a) (k : Fin b) => y (σ k)) x i j = if i = σ j then 1 else 0

Reindex Jacobian — sparse, hits 1 only at i = σ(j). Subsumes pdiv_id (set a = b, σ = id). Covers transpose, flatten, unflatten, slicing, any permutation.

theorem Proofs.pdiv_mul {m n : ℕ} (f g : Vec m → Vec n) (x : Vec m) (hf : DifferentiableAt ℝ f x) (hg : DifferentiableAt ℝ g x) (i : Fin m) (j : Fin n) : pdiv (fun (y : Vec m) (k : Fin n) => f y k * g y k) x i j = pdiv f x i j * g x j + f x j * pdiv g x i j

Product rule for pdiv. Vec n is a normed algebra over ℝ via Pi.normedAlgebra, so fderiv_mul applies directly to the pointwise product f * g. Requires both factors to be DifferentiableAt x.

theorem Proofs.pdiv_add {m n : ℕ} (f g : Vec m → Vec n) (x : Vec m) (hf : DifferentiableAt ℝ f x) (hg : DifferentiableAt ℝ g x) (i : Fin m) (j : Fin n) : pdiv (fun (y : Vec m) (k : Fin n) => f y k + g y k) x i j = pdiv f x i j + pdiv g x i j

Sum rule for pdiv. Requires both summands to be DifferentiableAt x.

theorem Proofs.pdiv_comp {m n p : ℕ} (f : Vec m → Vec n) (g : Vec n → Vec p) (x : Vec m) (hf : DifferentiableAt ℝ f x) (hg : DifferentiableAt ℝ g (f x)) (i : Fin m) (k : Fin p) : pdiv (g ∘ f) x i k = ∑ j : Fin n, pdiv f x i j * pdiv g (f x) j k

Chain rule for pdiv. Requires f differentiable at x and g differentiable at f x.

theorem Proofs.pdiv_const_smul {m n : ℕ} (c : ℝ) (f : Vec m → Vec n) (x : Vec m) (hf : DifferentiableAt ℝ f x) (i : Fin m) (j : Fin n) : pdiv (fun (y : Vec m) (k : Fin n) => c * f y k) x i j = c * pdiv f x i j

Scalar multiple rule for pdiv — pdiv_mul at a constant factor.

theorem Proofs.pdiv_coordFun {K : ℕ} (f : ℝ → ℝ) (f' : ℝ) (k : Fin K) (z : Vec K) (hf : HasDerivAt f f' (z k)) (j : Fin K) : pdiv (fun (z' : Vec K) (x : Fin 1) => f (z' k)) z j 0 = if j = k then f' else 0

A scalar function of ONE coordinate, lifted to Vec K → Vec 1, and its pdiv. pdiv_sigmoid's proof at one coordinate and a general f; the shape every summand of a per-class loss has.

theorem Proofs.pdiv_finset_sum {m n : ℕ} {α : Type u_1} [DecidableEq α] (S : Finset α) (f : α → Vec m → Vec n) (x : Vec m) (hdiff : ∀ s ∈ S, DifferentiableAt ℝ (f s) x) (i : Fin m) (j : Fin n) : pdiv (fun (y : Vec m) (k : Fin n) => ∑ s ∈ S, f s y k) x i j = ∑ s ∈ S, pdiv (f s) x i j

Finset-sum rule — derived from pdiv_add and pdiv_const by induction on the Finset. Linearity of the derivative extended to arbitrary finite sums. Requires each f s to be differentiable at x.

structure Proofs.HasVJP {m n : ℕ} (f : Vec m → Vec n) : Type

• backward : Vec m → Vec n → Vec m
• correct (x : Vec m) (dy : Vec n) (i : Fin m) : self.backward x dy i = ∑ j : Fin n, pdiv f x i j * dy j

noncomputable def Proofs.vjp_comp {m n p : ℕ} (f : Vec m → Vec n) (g : Vec n → Vec p) (hf_diff : Differentiable ℝ f) (hg_diff : Differentiable ℝ g) (hf : HasVJP f) (hg : HasVJP g) : HasVJP (g ∘ f)

Chain rule for VJPs — proved, no sorry. Requires f and g to be differentiable everywhere.

Equations

• Proofs.vjp_comp f g hf_diff hg_diff hf hg = { backward := fun (x : Proofs.Vec m) (dy : Proofs.Vec p) => hf.backward x (hg.backward (f x) dy), correct := ⋯ }

@[reducible]

noncomputable def Proofs.biPath {m n : ℕ} (f g : Vec m → Vec n) : Vec m → Vec n

Additive fan-in — proved, no sorry. Requires f and g to be differentiable everywhere.

Equations

• Proofs.biPath f g x i = f x i + g x i

noncomputable def Proofs.biPath_has_vjp {m n : ℕ} (f g : Vec m → Vec n) (hf_diff : Differentiable ℝ f) (hg_diff : Differentiable ℝ g) (hf : HasVJP f) (hg : HasVJP g) : HasVJP (biPath f g)

Equations

• Proofs.biPath_has_vjp f g hf_diff hg_diff hf hg = { backward := fun (x : Proofs.Vec m) (dy : Proofs.Vec n) (i : Fin m) => hf.backward x dy i + hg.backward x dy i, correct := ⋯ }

@[reducible]

noncomputable def Proofs.elemwiseProduct {n : ℕ} (f g : Vec n → Vec n) : Vec n → Vec n

Multiplicative fan-in — proved, no sorry. Requires f and g to be differentiable everywhere.

Equations

• Proofs.elemwiseProduct f g x i = f x i * g x i

noncomputable def Proofs.elemwiseProduct_has_vjp {n : ℕ} (f g : Vec n → Vec n) (hf_diff : Differentiable ℝ f) (hg_diff : Differentiable ℝ g) (hf : HasVJP f) (hg : HasVJP g) : HasVJP (elemwiseProduct f g)

Equations

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

def Proofs.identity_has_vjp (n : ℕ) : HasVJP fun (x : Vec n) => x

Identity VJP — proved, no sorry.

Equations

• Proofs.identity_has_vjp n = { backward := fun (_x dy : Proofs.Vec n) => dy, correct := ⋯ }

Why a separate HasVJPAt. The global HasVJP framework delivers a single backward function that's correct at every input. For non-smooth operators (relu, maxPool2, …) the only honest correct witness is the canonical pdiv-derived sum, which gives a trivially-rfl-true contract that doesn't pin down behavior at the kinks. HasVJPAt f x carries the same contract but only at a chosen smooth point x — exactly enough to discharge the chain rule under DifferentiableAt and to plug in real per-operator Jacobian formulas (pdiv_relu, pdiv3_maxPool2_smooth, …) instead of correct := rfl.

Smooth operators (dense, add, mul, softmax, batchNorm, …) keep their global HasVJP instances; we trivially lift to HasVJPAt at any point via HasVJP.toHasVJPAt when composing.

structure Proofs.HasVJPAt {m n : ℕ} (f : Vec m → Vec n) (x : Vec m) : Type

• backward : Vec n → Vec m
• correct (dy : Vec n) (i : Fin m) : self.backward dy i = ∑ j : Fin n, pdiv f x i j * dy j

def Proofs.HasVJP.toHasVJPAt {m n : ℕ} {f : Vec m → Vec n} (hf : HasVJP f) (x : Vec m) : HasVJPAt f x

Trivial lift: a global HasVJP gives a HasVJPAt at any point.

Equations

• hf.toHasVJPAt x = { backward := fun (dy : Proofs.Vec n) => hf.backward x dy, correct := ⋯ }

def Proofs.identity_has_vjp_at (n : ℕ) (x : Vec n) : HasVJPAt (fun (y : Vec n) => y) x

Identity pointwise VJP — trivial.

Equations

• Proofs.identity_has_vjp_at n x = (Proofs.identity_has_vjp n).toHasVJPAt x

noncomputable def Proofs.vjp_comp_at {m n p : ℕ} (f : Vec m → Vec n) (g : Vec n → Vec p) (x : Vec m) (hf_diff : DifferentiableAt ℝ f x) (hg_diff : DifferentiableAt ℝ g (f x)) (hf : HasVJPAt f x) (hg : HasVJPAt g (f x)) : HasVJPAt (g ∘ f) x

Chain rule for pointwise VJPs. Same shape as vjp_comp, but only requires DifferentiableAt at the relevant points (not everywhere). The pointwise analogue is what lets us compose through relu at smooth inputs.

Equations

• Proofs.vjp_comp_at f g x hf_diff hg_diff hf hg = { backward := fun (dy : Proofs.Vec p) => hf.backward (hg.backward dy), correct := ⋯ }

Mat m n and Vec (m * n) are in bijection by row-major flattening. This bijection lets us define pdivMat in terms of pdiv rather than introducing parallel axioms, and so derive the rank-2 chain, sum, and identity rules as theorems. The 5 local Jacobian theorems (matmul, scalarScale, transpose, rowIndep) are likewise derived from foundation rules — they state genuine calculus facts about specific operations, not structural framework.

noncomputable def Proofs.Mat.flatten {m n : ℕ} (A : Mat m n) : Vec (m * n)

Row-major flatten: Mat m n → Vec (m * n). Uses Mathlib's finProdFinEquiv : Fin m × Fin n ≃ Fin (m * n).

Equations

• A.flatten k = A (finProdFinEquiv.symm k).1 (finProdFinEquiv.symm k).2

noncomputable def Proofs.Mat.unflatten {m n : ℕ} (v : Vec (m * n)) : Mat m n

Row-major unflatten: Vec (m * n) → Mat m n.

Equations

• Proofs.Mat.unflatten v i j = v (finProdFinEquiv (i, j))

theorem Proofs.Mat.unflatten_flatten {m n : ℕ} (A : Mat m n) : unflatten A.flatten = A

Unflatten is a left inverse of flatten.

theorem Proofs.Mat.flatten_unflatten {m n : ℕ} (v : Vec (m * n)) : (unflatten v).flatten = v

Flatten is a left inverse of unflatten.

noncomputable def Proofs.pdivMat {a b c d : ℕ} (f : Mat a b → Mat c d) (A : Mat a b) (i : Fin a) (j : Fin b) (k : Fin c) (l : Fin d) : ℝ

Matrix partial derivative, defined in terms of pdiv on the row-major flattened Vec form. No longer an axiom — the rank-2 structural rules (chain/sum/id) now follow as theorems.

Equations

• Proofs.pdivMat f A i j k l = Proofs.pdiv (fun (v : Proofs.Vec (a * b)) => (f (Proofs.Mat.unflatten v)).flatten) A.flatten (finProdFinEquiv (i, j)) (finProdFinEquiv (k, l))

theorem Proofs.pdivMat_comp {a b c d e f : ℕ} (F : Mat a b → Mat c d) (G : Mat c d → Mat e f) (A : Mat a b) (hF_diff : DifferentiableAt ℝ (fun (v : Vec (a * b)) => (F (Mat.unflatten v)).flatten) A.flatten) (hG_diff : DifferentiableAt ℝ (fun (u : Vec (c * d)) => (G (Mat.unflatten u)).flatten) (F A).flatten) (i : Fin a) (j : Fin b) (k : Fin e) (l : Fin f) : pdivMat (G ∘ F) A i j k l = ∑ p : Fin c, ∑ q : Fin d, pdivMat F A i j p q * pdivMat G (F A) p q k l

Chain rule for pdivMat — now a theorem, derived from pdiv_comp via the row-major flatten bijection.

theorem Proofs.pdivMat_add {a b c d : ℕ} (F G : Mat a b → Mat c d) (A : Mat a b) (hF_diff : DifferentiableAt ℝ (fun (v : Vec (a * b)) => (F (Mat.unflatten v)).flatten) A.flatten) (hG_diff : DifferentiableAt ℝ (fun (v : Vec (a * b)) => (G (Mat.unflatten v)).flatten) A.flatten) (i : Fin a) (j : Fin b) (k : Fin c) (l : Fin d) : pdivMat (fun (M : Mat a b) (r : Fin c) (s : Fin d) => F M r s + G M r s) A i j k l = pdivMat F A i j k l + pdivMat G A i j k l

Sum rule for pdivMat — theorem, via pdiv_add. Requires both flattened summands to be differentiable at flatten A.

theorem Proofs.pdivMat_id {a b : ℕ} (A : Mat a b) (i : Fin a) (j : Fin b) (k : Fin a) (l : Fin b) : pdivMat (fun (M : Mat a b) => M) A i j k l = if i = k ∧ j = l then 1 else 0

Identity Jacobian for pdivMat — theorem, via pdiv_id.

structure Proofs.HasVJPMat {a b c d : ℕ} (f : Mat a b → Mat c d) : Type

Matrix-level VJP: given a matrix-valued function of a matrix, a correct backward function contracts the pdivMat Jacobian against the output cotangent. Mirrors HasVJP for Vec.

• backward : Mat a b → Mat c d → Mat a b
• correct (A : Mat a b) (dY : Mat c d) (i : Fin a) (j : Fin b) : self.backward A dY i j = ∑ k : Fin c, ∑ l : Fin d, pdivMat f A i j k l * dY k l

noncomputable def Proofs.vjpMat_comp {a b c d e f : ℕ} (F : Mat a b → Mat c d) (G : Mat c d → Mat e f) (hF_diff : Differentiable ℝ fun (v : Vec (a * b)) => (F (Mat.unflatten v)).flatten) (hG_diff : Differentiable ℝ fun (u : Vec (c * d)) => (G (Mat.unflatten u)).flatten) (hF : HasVJPMat F) (hG : HasVJPMat G) : HasVJPMat (G ∘ F)

Chain rule for matrix VJPs — proved, no sorry. Direct transcription of vjp_comp to rank-2 indices.

Equations

• Proofs.vjpMat_comp F G hF_diff hG_diff hF hG = { backward := fun (A : Proofs.Mat a b) (dY : Proofs.Mat e f) => hF.backward A (hG.backward (F A) dY), correct := ⋯ }

@[reducible]

noncomputable def Proofs.biPathMat {a b c d : ℕ} (F G : Mat a b → Mat c d) : Mat a b → Mat c d

Additive fan-in for matrices — proved, no sorry.

Equations

• Proofs.biPathMat F G M r s = F M r s + G M r s

noncomputable def Proofs.biPathMat_has_vjp {a b c d : ℕ} (F G : Mat a b → Mat c d) (hF_diff : Differentiable ℝ fun (v : Vec (a * b)) => (F (Mat.unflatten v)).flatten) (hG_diff : Differentiable ℝ fun (v : Vec (a * b)) => (G (Mat.unflatten v)).flatten) (hF : HasVJPMat F) (hG : HasVJPMat G) : HasVJPMat (biPathMat F G)

Equations

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

noncomputable def Proofs.identityMat_has_vjp (a b : ℕ) : HasVJPMat fun (M : Mat a b) => M

Identity VJP for matrices — proved, no sorry.

Equations

• Proofs.identityMat_has_vjp a b = { backward := fun (_A dY : Proofs.Mat a b) => dY, correct := ⋯ }

noncomputable def Proofs.hasVJPMat_to_hasVJP {a b c d : ℕ} {f : Mat a b → Mat c d} (hf : HasVJPMat f) : HasVJP fun (v : Vec (a * b)) => (f (Mat.unflatten v)).flatten

Bridge: HasVJPMat → HasVJP via the Mat.flatten bijection.

Given a matrix-level VJP for f : Mat a b → Mat c d, produce a vector-level VJP for the flattened version fun v : Vec (a*b) => Mat.flatten (f (Mat.unflatten v)). The backward reshapes the input/output flat vectors to matrices, applies the matrix backward, and flattens the result.

Lets us compose HasVJPMat pieces (vit_body, transformer blocks) with rank-crossing pieces (patch embed, classifier head) that live natively as Vec → Vec by first bridging everything to HasVJP.

Equations

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

The three theorems here are local Jacobians for the operations that appear in scaled dot-product attention's backward pass:

1. pdivMat_matmul_left_const — right-factor varies, left factor fixed: ∂(C · B')_{kl} / ∂B'_{ij} = C_{ki} · [l = j].
2. pdivMat_matmul_right_const — left factor varies, right factor fixed: ∂(A' · D)_{kl} / ∂A'_{ij} = D_{jl} · [i = k].
3. pdivMat_rowIndep — functions that act row-wise have block-diagonal Jacobians, with the per-row block equal to the vector Jacobian of the row function g.

Each is a direct transcription of an elementary calculus fact. They are numerically gradient-checked in check_jacobians.py.

theorem Proofs.pdivMat_matmul_left_const {m p q : ℕ} (C : Mat m p) (B : Mat p q) (i : Fin p) (j : Fin q) (k : Fin m) (l : Fin q) : pdivMat (fun (B' : Mat p q) => C.mul B') B i j k l = if l = j then C k i else 0

Matmul Jacobian (left-const) — theorem, derived from pdiv_finset_sum + pdiv_mul + pdiv_const + pdiv_reindex.

theorem Proofs.pdivMat_matmul_right_const {m p q : ℕ} (A : Mat m p) (D : Mat p q) (i : Fin m) (j : Fin p) (k : Fin m) (l : Fin q) : pdivMat (fun (A' : Mat m p) => A'.mul D) A i j k l = if i = k then D j l else 0

Matmul Jacobian (right-const) — theorem, same recipe as the left-const case with roles swapped.

theorem Proofs.pdivMat_rowIndep {m n p : ℕ} (g : Vec n → Vec p) (h_g_diff : Differentiable ℝ g) (A : Mat m n) (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 — proved (planning/archive/VJP.md follow-up D).

For a row-independent function M ↦ (r ↦ g (M r)), the (i,j,k,l) Jacobian entry is pdiv g (A i) j l when i = k and 0 otherwise.

Requires Differentiable ℝ g: without it, the flattened Pi-valued function may be non-differentiable at Mat.flatten A (per differentiable_pi's coordinate-wise condition), making fderiv = 0 junk and breaking the per-row decomposition.

noncomputable def Proofs.rowwise_has_vjp_mat {m n p : ℕ} {g : Vec n → Vec p} (hg : HasVJP g) (hg_diff : Differentiable ℝ g) : HasVJPMat fun (A : Mat m n) (r : Fin m) => g (A r)

Row-wise lifting of a HasVJP (Phase 8, Tensor-level).

Given any g : Vec n → Vec p with a proved HasVJP, applying g independently to each row of a matrix A : Mat m n gives a HasVJPMat on Mat m n → Mat m p. The backward is just g.backward applied per row. Generalizes rowSoftmax_has_vjp_mat: any per-token operation (LayerNorm, GELU, dense, activation) lifts to a per-sequence matrix operation via this one helper.

Equations

• Proofs.rowwise_has_vjp_mat hg hg_diff = { backward := fun (A : Proofs.Mat m n) (dY : Proofs.Mat m p) (r : Fin m) (c : Fin n) => hg.backward (A r) (dY r) c, correct := ⋯ }

Per-head / per-slab column independence

Multi-head attention applies the same per-head function to each of heads column slabs of width d_in from a Mat n (heads * d_in) input. The column-slab analog of rowwise_has_vjp_mat factors that vmap-over-heads structure: each head's output depends only on its own slab of the input, so the matrix Jacobian is block-diagonal across the head axis.

noncomputable def Proofs.colSlabApply {n heads d_in d_out : ℕ} (g : Mat n d_in → Mat n d_out) : Mat n (heads * d_in) → Mat n (heads * d_out)

Apply g : Mat n d_in → Mat n d_out to each of the heads column slabs of width d_in in a Mat n (heads * d_in) input, producing a Mat n (heads * d_out) output. Output column (h, j_out) is column j_out of g (slab h M), where slab h M extracts the d_in-wide column block at head index h.

Equations

• Proofs.colSlabApply g M r hj = g (fun (r' : Fin n) (j_in : Fin d_in) => M r' (finProdFinEquiv ((finProdFinEquiv.symm hj).1, j_in))) r (finProdFinEquiv.symm hj).2

theorem Proofs.pdivMat_colIndep {n heads d_in d_out : ℕ} (g : Mat n d_in → Mat n d_out) (h_g_diff : Differentiable ℝ fun (v : Vec (n * d_in)) => (g (Mat.unflatten v)).flatten) (A : Mat n (heads * d_in)) (i : Fin n) (h_j : Fin heads) (j' : Fin d_in) (k : Fin n) (h_l : Fin heads) (j'' : Fin d_out) : pdivMat (colSlabApply g) A i (finProdFinEquiv (h_j, j')) k (finProdFinEquiv (h_l, j'')) = if h_j = h_l then pdivMat g (fun (r' : Fin n) (j_in : Fin d_in) => A r' (finProdFinEquiv (h_l, j_in))) i j' k j'' else 0

Column-slab independence Jacobian — column-axis analog of pdivMat_rowIndep. For a slab-applied function colSlabApply g, the Jacobian is block-diagonal across the heads axis: zero unless the input slab h_j matches the output slab h_l, otherwise equal to pdivMat g on that slab.

Requires Differentiable ℝ (flat g) for the same reason as pdivMat_rowIndep: the Pi-valued flat form must be differentiable everywhere so fderiv doesn't fall back to junk-default 0.

noncomputable def Proofs.colSlabwise_has_vjp_mat {n heads d_in d_out : ℕ} {g : Mat n d_in → Mat n d_out} (hg : HasVJPMat g) (hg_diff : Differentiable ℝ fun (v : Vec (n * d_in)) => (g (Mat.unflatten v)).flatten) : HasVJPMat (colSlabApply g)

Lift HasVJPMat g to column-slab vmap — column-axis analog of rowwise_has_vjp_mat. Given g : Mat n d_in → Mat n d_out with a matrix VJP, applying g independently to each of heads-many column slabs gives a HasVJPMat for colSlabApply g. The backward applies g.backward per slab.

Equations

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

Ternary matrix VJP

For ternary-input functions like SDPA (Q, K, V) ↦ out, package the three per-input VJPs as a single structure analogous to HasVJPMat. The backward returns the triple of per-input gradients; correctness holds independently for each input (with the others fixed).

structure Proofs.HasVJPMat3 {n d_in d_out : ℕ} (F : Mat n d_in → Mat n d_in → Mat n d_in → Mat n d_out) : Type

VJP structure for Mat × Mat × Mat → Mat functions where all three inputs share the same shape Mat n d_in and the output is Mat n d_out. Backward returns the triple of per-input gradients; correct_{1,2,3} ensure each gradient matches the partial derivative treating the other two inputs as constants.

• backward : Mat n d_in → Mat n d_in → Mat n d_in → Mat n d_out → Mat n d_in × Mat n d_in × Mat n d_in
• correct_1 (A B C : Mat n d_in) (dY : Mat n d_out) (i : Fin n) (j : Fin d_in) : (self.backward A B C dY).1 i j = ∑ k : Fin n, ∑ l : Fin d_out, pdivMat (fun (A' : Mat n d_in) => F A' B C) A i j k l * dY k l
• correct_2 (A B C : Mat n d_in) (dY : Mat n d_out) (i : Fin n) (j : Fin d_in) : (self.backward A B C dY).2.1 i j = ∑ k : Fin n, ∑ l : Fin d_out, pdivMat (fun (B' : Mat n d_in) => F A B' C) B i j k l * dY k l
• correct_3 (A B C : Mat n d_in) (dY : Mat n d_out) (i : Fin n) (j : Fin d_in) : (self.backward A B C dY).2.2 i j = ∑ k : Fin n, ∑ l : Fin d_out, pdivMat (fun (C' : Mat n d_in) => F A B C') C i j k l * dY k l

theorem Proofs.pdivMat_scalarScale {m n : ℕ} (s : ℝ) (A : Mat m n) (i : Fin m) (j : Fin n) (k : Fin m) (l : Fin n) : pdivMat (fun (M : Mat m n) (r : Fin m) (c : Fin n) => s * M r c) A i j k l = if i = k ∧ j = l then s else 0

Scalar-scale Jacobian — theorem, derived from pdiv_mul + pdiv_const + pdiv_id via the flatten bijection. ∂(s · A')_{kl} / ∂A'_{ij} = s · δ_{ik,jl}.

theorem Proofs.pdivMat_transpose {m n : ℕ} (A : Mat m n) (i : Fin m) (j k : Fin n) (l : Fin m) : pdivMat (fun (M : Mat m n) => M.transpose) A i j k l = if j = k ∧ i = l then 1 else 0

Transpose Jacobian — theorem, derived from pdiv_reindex via the flatten bijection. ∂A^T_{kl} / ∂A_{ij} = δ_{l=i, k=j}.

noncomputable def Proofs.matmul_left_const_has_vjp {m p q : ℕ} (C : Mat m p) : HasVJPMat fun (B' : Mat p q) => C.mul B'

Matmul with right factor varying, left factor fixed — proved.

f : Mat p q → Mat m q, f B' = C · B'. Backward: dB' = C^T · dY.

Equations

• Proofs.matmul_left_const_has_vjp C = { backward := fun (_B : Proofs.Mat p q) (dY : Proofs.Mat m q) (i : Fin p) (j : Fin q) => ∑ k : Fin m, C k i * dY k j, correct := ⋯ }

noncomputable def Proofs.matmul_right_const_has_vjp {m p q : ℕ} (D : Mat p q) : HasVJPMat fun (A' : Mat m p) => A'.mul D

Matmul with left factor varying, right factor fixed — proved.

f : Mat m p → Mat m q, f A' = A' · D. Backward: dA' = dY · D^T.

Equations

• Proofs.matmul_right_const_has_vjp D = { backward := fun (_A : Proofs.Mat m p) (dY : Proofs.Mat m q) (i : Fin m) (j : Fin p) => ∑ l : Fin q, dY i l * D j l, correct := ⋯ }

noncomputable def Proofs.scalarScale_has_vjp {m n : ℕ} (s : ℝ) : HasVJPMat fun (M : Mat m n) (r : Fin m) (c : Fin n) => s * M r c

Scalar-scale VJP — proved. Backward: dA = s · dY.

Equations

• Proofs.scalarScale_has_vjp s = { backward := fun (_A dY : Proofs.Mat m n) (i : Fin m) (j : Fin n) => s * dY i j, correct := ⋯ }

noncomputable def Proofs.transpose_has_vjp {m n : ℕ} : HasVJPMat fun (M : Mat m n) => M.transpose

Transpose VJP — proved. Backward: dA = (dY)^T.

Equations

• Proofs.transpose_has_vjp = { backward := fun (_A : Proofs.Mat m n) (dY : Proofs.Mat n m) (i : Fin m) (j : Fin n) => dY j i, correct := ⋯ }

@[reducible, inline]

abbrev Proofs.Tensor3 (c h w : ℕ) : Type

A 3D feature map: channels × height × width (single sample).

Equations

• Proofs.Tensor3 c h w = (Fin c → Fin h → Fin w → ℝ)

noncomputable def Proofs.Tensor3.flatten {c h w : ℕ} (T : Tensor3 c h w) : Vec (c * h * w)

Row-major flatten: Tensor3 c h w → Vec (c * h * w). Two nested finProdFinEquiv calls: first bundle (ci, hi) into Fin (c*h), then bundle with wi into Fin ((c*h)*w) = Fin (c*h*w).

Equations

• T.flatten k = T (finProdFinEquiv.symm (finProdFinEquiv.symm k).1).1 (finProdFinEquiv.symm (finProdFinEquiv.symm k).1).2 (finProdFinEquiv.symm k).2

noncomputable def Proofs.Tensor3.unflatten {c h w : ℕ} (v : Vec (c * h * w)) : Tensor3 c h w

Row-major unflatten: inverse of flatten.

Equations

• Proofs.Tensor3.unflatten v ci hi wi = v (finProdFinEquiv (finProdFinEquiv (ci, hi), wi))

theorem Proofs.Tensor3.unflatten_flatten {c h w : ℕ} (T : Tensor3 c h w) : unflatten T.flatten = T

theorem Proofs.Tensor3.flatten_unflatten {c h w : ℕ} (v : Vec (c * h * w)) : (unflatten v).flatten = v

theorem Proofs.Tensor3.flatten_differentiable {c h w : ℕ} : Differentiable ℝ flatten

Tensor3.flatten is differentiable. It is a coordinate reindexing: each output coordinate flatten x k is the single input coordinate x (decode k), hence a projection. differentiable_pi reduces to per-coordinate differentiability, which fun_prop discharges via the eval/projection rule.

theorem Proofs.Tensor3.unflatten_differentiable {c h w : ℕ} : Differentiable ℝ unflatten

Tensor3.unflatten is differentiable. The inverse reindexing: each output coordinate unflatten v ci hi wi is the single input coordinate v (encode (ci,hi,wi)). Three nested differentiable_pi peel the Fin c → Fin h → Fin w → ℝ pi structure; fun_prop discharges the innermost projection.

noncomputable def Proofs.pdiv3 {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} (f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂) (x : Tensor3 c₁ h₁ w₁) (ci : Fin c₁) (hi : Fin h₁) (wi : Fin w₁) (co : Fin c₂) (ho : Fin h₂) (wo : Fin w₂) : ℝ

3D partial derivative — now a definition via the triple-nested flatten bijection, no longer an axiom. The four structural rules (comp / add / id) follow as theorems. Operator-specific VJPs at rank 3 (conv2d_has_vjp3, maxPool2_has_vjp3, depthwise_has_vjp3) are bundled HasVJP3 defs in their respective files — those state specific Jacobian values, not framework.

Equations

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

theorem Proofs.pdiv3_comp {c₁ h₁ w₁ c₂ h₂ w₂ c₃ h₃ w₃ : ℕ} (f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂) (g : Tensor3 c₂ h₂ w₂ → Tensor3 c₃ h₃ w₃) (x : Tensor3 c₁ h₁ w₁) (hf_diff : DifferentiableAt ℝ (fun (v : Vec (c₁ * h₁ * w₁)) => (f (Tensor3.unflatten v)).flatten) x.flatten) (hg_diff : DifferentiableAt ℝ (fun (u : Vec (c₂ * h₂ * w₂)) => (g (Tensor3.unflatten u)).flatten) (f x).flatten) (ci : Fin c₁) (hi : Fin h₁) (wi : Fin w₁) (ck : Fin c₃) (hk : Fin h₃) (wk : Fin w₃) : pdiv3 (g ∘ f) x ci hi wi ck hk wk = ∑ cj : Fin c₂, ∑ hj : Fin h₂, ∑ wj : Fin w₂, pdiv3 f x ci hi wi cj hj wj * pdiv3 g (f x) cj hj wj ck hk wk

Chain rule for 3D partial derivatives — theorem, via pdiv_comp and two applications of Fintype.sum_equiv + sum_prod_type. Requires the flattened forms of f and g to be differentiable at the relevant points.

structure Proofs.HasVJP3 {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} (f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂) : Type

VJP for 3D→3D functions.

• backward : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂ → Tensor3 c₁ h₁ w₁
• correct (x : Tensor3 c₁ h₁ w₁) (dy : Tensor3 c₂ h₂ w₂) (ci : Fin c₁) (hi : Fin h₁) (wi : Fin w₁) : self.backward x dy ci hi wi = ∑ co : Fin c₂, ∑ ho : Fin h₂, ∑ wo : Fin w₂, pdiv3 f x ci hi wi co ho wo * dy co ho wo

noncomputable def Proofs.vjp3_comp {c₁ h₁ w₁ c₂ h₂ w₂ c₃ h₃ w₃ : ℕ} (f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂) (g : Tensor3 c₂ h₂ w₂ → Tensor3 c₃ h₃ w₃) (hf_diff : Differentiable ℝ fun (v : Vec (c₁ * h₁ * w₁)) => (f (Tensor3.unflatten v)).flatten) (hg_diff : Differentiable ℝ fun (u : Vec (c₂ * h₂ * w₂)) => (g (Tensor3.unflatten u)).flatten) (hf : HasVJP3 f) (hg : HasVJP3 g) : HasVJP3 (g ∘ f)

Chain rule for 3D VJPs — proved, no sorry. Requires the flattened forms of f and g to be differentiable everywhere.

Equations

• Proofs.vjp3_comp f g hf_diff hg_diff hf hg = { backward := fun (x : Proofs.Tensor3 c₁ h₁ w₁) (dy : Proofs.Tensor3 c₃ h₃ w₃) => hf.backward x (hg.backward (f x) dy), correct := ⋯ }

structure Proofs.HasVJPAt3 {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} (f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂) (x : Tensor3 c₁ h₁ w₁) : Type

Tensor3 analogue of HasVJPAt: the same pdiv3-sum contract, but only required at the chosen smooth point x. The natural home for maxPool2_has_vjp_at3 and any other kinked Tensor3 operator.

• backward : Tensor3 c₂ h₂ w₂ → Tensor3 c₁ h₁ w₁
• correct (dy : Tensor3 c₂ h₂ w₂) (ci : Fin c₁) (hi : Fin h₁) (wi : Fin w₁) : self.backward dy ci hi wi = ∑ co : Fin c₂, ∑ ho : Fin h₂, ∑ wo : Fin w₂, pdiv3 f x ci hi wi co ho wo * dy co ho wo

def Proofs.HasVJP3.toHasVJPAt3 {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} {f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂} (hf : HasVJP3 f) (x : Tensor3 c₁ h₁ w₁) : HasVJPAt3 f x

Trivial lift: a global HasVJP3 gives a HasVJPAt3 at any point.

Equations

• hf.toHasVJPAt3 x = { backward := fun (dy : Proofs.Tensor3 c₂ h₂ w₂) => hf.backward x dy, correct := ⋯ }

noncomputable def Proofs.vjp3_comp_at {c₁ h₁ w₁ c₂ h₂ w₂ c₃ h₃ w₃ : ℕ} (f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂) (g : Tensor3 c₂ h₂ w₂ → Tensor3 c₃ h₃ w₃) (x : Tensor3 c₁ h₁ w₁) (hf_diff : DifferentiableAt ℝ (fun (v : Vec (c₁ * h₁ * w₁)) => (f (Tensor3.unflatten v)).flatten) x.flatten) (hg_diff : DifferentiableAt ℝ (fun (u : Vec (c₂ * h₂ * w₂)) => (g (Tensor3.unflatten u)).flatten) (f x).flatten) (hf : HasVJPAt3 f x) (hg : HasVJPAt3 g (f x)) : HasVJPAt3 (g ∘ f) x

Chain rule for pointwise Tensor3 VJPs. Tensor3 analogue of vjp_comp_at. Requires DifferentiableAt only at the relevant points (on the flattened forms), which is what vjp3_comp_at consumers built from _at instances of kinked operators can actually supply.

Equations

• Proofs.vjp3_comp_at f g x hf_diff hg_diff hf hg = { backward := fun (dy : Proofs.Tensor3 c₃ h₃ w₃) => hf.backward (hg.backward dy), correct := ⋯ }

noncomputable def Proofs.hasVJP3_to_hasVJP {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} {f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂} (hf : HasVJP3 f) : HasVJP fun (v : Vec (c₁ * h₁ * w₁)) => (f (Tensor3.unflatten v)).flatten

Bridge: HasVJP3 → HasVJP via the Tensor3.flatten bijection.

Rank-3 analogue of hasVJPMat_to_hasVJP. Given a Tensor3-level VJP for f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂, produce a vector-level VJP for the flattened fun v => Tensor3.flatten (f (Tensor3.unflatten v)). The backward decodes the flat index in two finProdFinEquiv.symm levels (matching pdiv3's row-major encode), applies the Tensor3 backward, and the closing collapse folds the triple co/ho/wo sum back to the single flat sum via two Fintype.sum_prod_type + Fintype.sum_equiv finProdFinEquiv reindexes.

Equations

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

noncomputable def Proofs.hasVJPAt3_to_hasVJPAt {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} {f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂} {x : Tensor3 c₁ h₁ w₁} (hf : HasVJPAt3 f x) : HasVJPAt (fun (v : Vec (c₁ * h₁ * w₁)) => (f (Tensor3.unflatten v)).flatten) x.flatten

Bridge: HasVJPAt3 → HasVJPAt via the Tensor3.flatten bijection.

Smooth-point analogue of hasVJP3_to_hasVJP, with x fixed. Needed for kinked operators (e.g. maxPool2) that only carry HasVJPAt3. Same two-level index decode and triple→flat reindex collapse.

Equations

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

theorem Proofs.pdiv3_id {c h w : ℕ} (x : Tensor3 c h w) (ci : Fin c) (hi : Fin h) (wi : Fin w) (co : Fin c) (ho : Fin h) (wo : Fin w) : pdiv3 (fun (t : Tensor3 c h w) => t) x ci hi wi co ho wo = if ci = co ∧ hi = ho ∧ wi = wo then 1 else 0

Identity Jacobian for Tensor3 — theorem, via pdiv_id and injectivity of the nested finProdFinEquiv.

def Proofs.identity3_has_vjp (c h w : ℕ) : HasVJP3 fun (x : Tensor3 c h w) => x

Equations

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

theorem Proofs.pdiv3_add {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} (f g : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂) (x : Tensor3 c₁ h₁ w₁) (hf_diff : DifferentiableAt ℝ (fun (v : Vec (c₁ * h₁ * w₁)) => (f (Tensor3.unflatten v)).flatten) x.flatten) (hg_diff : DifferentiableAt ℝ (fun (v : Vec (c₁ * h₁ * w₁)) => (g (Tensor3.unflatten v)).flatten) x.flatten) (ci : Fin c₁) (hi : Fin h₁) (wi : Fin w₁) (co : Fin c₂) (ho : Fin h₂) (wo : Fin w₂) : pdiv3 (fun (y : Tensor3 c₁ h₁ w₁) (c : Fin c₂) (h : Fin h₂) (w : Fin w₂) => f y c h w + g y c h w) x ci hi wi co ho wo = pdiv3 f x ci hi wi co ho wo + pdiv3 g x ci hi wi co ho wo

Sum rule for Tensor3 partial derivatives — theorem, via pdiv_add.

@[reducible]

noncomputable def Proofs.biPath3 {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} (f g : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂) : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂

Equations

• Proofs.biPath3 f g x c h w = f x c h w + g x c h w

noncomputable def Proofs.biPath3_has_vjp {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} (f g : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂) (hf_diff : Differentiable ℝ fun (v : Vec (c₁ * h₁ * w₁)) => (f (Tensor3.unflatten v)).flatten) (hg_diff : Differentiable ℝ fun (v : Vec (c₁ * h₁ * w₁)) => (g (Tensor3.unflatten v)).flatten) (hf : HasVJP3 f) (hg : HasVJP3 g) : HasVJP3 (biPath3 f g)

Equations

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