Tensor algebra and the VJP framework — the root of the proof suite #

Vec n := Fin n → ℝ, Mat m n, Tensor3 c h w, and the backward-pass calculus over them, grounded in Mathlib's Fréchet derivative:

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

Every structural rule is a theorem against Mathlib's API; the bilinear ones carry Differentiable hypotheses that propagate through every downstream chapter.

Contents, in file order:

section	what
Types, Matrix Operations	`Vec`, `Mat`, `Mat.mulVec`/`mul`/`transpose`/`outer`, `basisVec`
Differentiation	`pdiv` and its rules (`pdiv_comp`, `pdiv_add`, `pdiv_mul`, `pdiv_reindex`, …); the linear-map kits `pdiv_of_affine` / `pdiv_of_linear`, `pdiv_elementwise`, `pdiv_finset_sum`, `pdiv_const_smul`
VJP Framework, Pointwise VJP	`HasVJP` (global) and `HasVJPAt` (at a point), each `@[ext]` and a `Subsingleton`, `canonical`, `backward_unique`, `vjpComp` / `vjpCompAt` / `vjpCompDiffAt` and their `_backward` peels, `biPath`
flattening	`Mat.flatten` / `unflatten` with `_apply` and `@[simp]` round-trips, `sum_finProdFinEquiv`
Matrix-level	`pdivMat`, `HasVJPMat`, `vjpMatComp`, row-independence (`pdivMat_rowIndep*`, `rowwiseHasVJPMat`), and the matmul / scale / transpose Jacobians and VJPs. The multi-head column-slab kit and `HasVJPMat3` live in `Attention.lean`
3D tensors	`Tensor3`, `pdiv3`, `HasVJP3` / `HasVJPAt3` and `HasVJP3.toHasVJP` — nets compose at `Vec` through it

Zero project axioms: everything closes under propext, Classical.choice, Quot.sound.

source

@[reducible, inline]

abbrev Proofs.Vec (n : ℕ) :

Type

Equations

Proofs.Vec n = (Fin n → ℝ)

Instances For

source

@[reducible, inline]

abbrev Proofs.Mat (m n : ℕ) :

Type

Equations

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

Instances For

source

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

Instances For

source

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

Instances For

source

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

Instances For

source

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

Instances For

source

@[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

Instances For

source

@[simp]

theorem Proofs.basisVec_apply {m : ℕ} (i j : Fin m) :

basisVec i j = if j = i then 1 else 0

source

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 := ⋯ }

Instances For

source

@[simp]

theorem Proofs.reindexCLM_apply {a b : ℕ} (σ : Fin b → Fin a) (y : Vec a) :

(reindexCLM σ) y = fun (k : Fin b) => y (σ k)

source

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

Instances For

source

theorem Proofs.pdiv_eq_of_hasFDerivAt {m n : ℕ} {F : Vec m → Vec n} {F' : Vec m →L[ℝ ] Vec n} {x : Vec m} (h : HasFDerivAt F F' x) (i : Fin m) (j : Fin n) :

pdiv F x i j = F' (basisVec i) j

A known derivative reads off as pdiv: output j of F' on the i-th basis vector.

source

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 — δᵢⱼ.

source

theorem Proofs.pdiv_eq_fderiv_coord {m n : ℕ} {f : Vec m → Vec n} {x : Vec m} (hf : DifferentiableAt ℝ f x) (i : Fin m) (j : Fin n) :

pdiv f x i j = (fderiv ℝ (fun (y : Vec m) => f y j) x) (basisVec i)

A pdiv entry is the derivative of one output coordinate — pdiv f x i j is fderiv f x (basisVec i) j, and at a differentiable point that is the i-th directional derivative of y ↦ f y j.

source

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.

source

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.

source

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.

source

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.

source

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.

source

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.

source

theorem Proofs.pdiv_elementwise {n : ℕ} (φ : ℝ → ℝ) (x : Vec n) (hφ : ∀ (k : Fin n), DifferentiableAt ℝ φ (x k)) (i j : Fin n) :

pdiv (fun (y : Vec n) (k : Fin n) => φ (y k)) x i j = if i = j then deriv φ (x i) else 0

Elementwise rule — a scalar function applied to every coordinate has a diagonal Jacobian, the scalar derivative on the diagonal: each coordinate is φ ∘ proj k (HasDerivAt.comp_hasFDerivAt), assembled by hasFDerivAt_pi. GELU, swish and sigmoid are this with φ their scalar function.

source

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. The elementwise rule at one coordinate and a general f; the shape every summand of a per-class loss has.

source

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 — linearity of the derivative extended to arbitrary finite sums (fderiv_fun_sum). Requires each f s to be differentiable at x.

source

theorem Proofs.pdiv_lift_sum {P : ℕ} {α : Type u_1} [DecidableEq α] (S : Finset α) (f : α → Vec P → ℝ) (x : Vec P) (h : ∀ k ∈ S, DifferentiableAt ℝ (fun (z : Vec P) (x : Fin 1) => f k z) x) (j : Fin P) :

pdiv (fun (z : Vec P) (x : Fin 1) => ∑ k ∈ S, f k z) x j 0 = ∑ k ∈ S, pdiv (fun (z : Vec P) (x : Fin 1) => f k z) x j 0

pdiv_finset_sum at a scalar loss: the gradient of a sum of scalar losses (lifted to Vec 1) is the sum of their gradients. The form every summed-loss cotangent takes.

source

theorem Proofs.pdiv_of_affine {m n : ℕ} (f : Vec m → Vec n) (c : Vec n) (hadd : ∀ (u v : Vec m), f (u + v) = f u + f v) (hsmul : ∀ (a : ℝ) (v : Vec m), f (a • v) = a • f v) (x : Vec m) (i : Fin m) (j : Fin n) :

pdiv (fun (v : Vec m) => f v + c) x i j = f (basisVec i) j

Affine rule — if f is additive and homogeneous it is a linear map on the finite-dimensional Vec m, hence continuous (LinearMap.continuous_on_pi), and the Jacobian of v ↦ f v + c at every point is f read on the basis vector. Every conv / depthwise / dense / patch-embed / pooling / BN-affine Jacobian in the suite is this with c the bias: two one-line linearity side goals instead of distributing pdiv through the sum.

source

theorem Proofs.pdiv_of_linear {m n : ℕ} (f : Vec m → Vec n) (hadd : ∀ (u v : Vec m), f (u + v) = f u + f v) (hsmul : ∀ (a : ℝ) (v : Vec m), f (a • v) = a • f v) (x : Vec m) (i : Fin m) (j : Fin n) :

pdiv f x i j = f (basisVec i) j

Linear rule, unbundled — pdiv_of_affine at c = 0.

source

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

Type

A vector–Jacobian product for f: a backward map with the proof that at every x it contracts f's Jacobian (pdiv) against the cotangent. correct is the whole content. HasVJP.canonical inhabits it for every f (its backward is the contraction itself), so a witness says something only when its backward is a formula proved equal to that contraction.

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

Instances For

source

theorem Proofs.HasVJP.ext {m n : ℕ} {f : Vec m → Vec n} {x y : HasVJP f} (backward : x.backward = y.backward) :

x = y

source

theorem Proofs.HasVJP.ext_iff {m n : ℕ} {f : Vec m → Vec n} {x y : HasVJP f} :

x = y ↔ x.backward = y.backward

source

instance Proofs.instSubsingletonHasVJP {m n : ℕ} {f : Vec m → Vec n} :

Subsingleton (HasVJP f)

correct pins backward at every input, so a map has at most one VJP witness.

source

noncomputable def Proofs.HasVJP.canonical {m n : ℕ} (f : Vec m → Vec n) :

HasVJP f

The canonical witness — the backward IS the pdiv contraction, so correct is rfl. Exists for every f; a hand-written backward is tied to it by HasVJP.backward_unique.

Equations

Proofs.HasVJP.canonical f = { backward := fun (x : Proofs.Vec m) (dy : Proofs.Vec n) (i : Fin m) => ∑ j : Fin n, Proofs.pdiv f x i j * dy j, correct := ⋯ }

Instances For

source

def Proofs.HasVJP.congr {m n : ℕ} {f g : Vec m → Vec n} (h : f = g) (hf : HasVJP f) :

HasVJP g

Transport a witness along an equality of maps. The backward is carried over verbatim, so unlike h ▸ hf it is never hidden behind a cast (congr_backward is rfl).

Equations

Proofs.HasVJP.congr h hf = { backward := hf.backward, correct := ⋯ }

Instances For

source

@[simp]

theorem Proofs.HasVJP.congr_backward {m n : ℕ} {f g : Vec m → Vec n} (h : f = g) (hf : HasVJP f) :

(congr h hf).backward = hf.backward

source

theorem Proofs.HasVJP.backward_unique_of_eq {m n : ℕ} {f g : Vec m → Vec n} (hfg : f = g) (h₁ : HasVJP f) (h₂ : HasVJP g) (x : Vec m) (dy : Vec n) :

h₁.backward x dy = h₂.backward x dy

Two VJP witnesses for EQUAL maps have the same backward. Both .correct to the same ∑ pdiv f x i j * dy j. Going through .correct rather than hfg ▸ · avoids an Eq.mpr-blocked backward when the witnesses have different types (a respelling of f).

source

theorem Proofs.HasVJP.backward_unique {m n : ℕ} {f : Vec m → Vec n} (h₁ h₂ : HasVJP f) (x : Vec m) (dy : Vec n) :

h₁.backward x dy = h₂.backward x dy

Any two VJP witnesses for the same map have the same backward — the backward is a property of f, not of how the witness was assembled. Lets a hand-written chain be tied to a tactic-built witness without unfolding it.

source

@[reducible, inline]

abbrev Proofs.IsHomog {a b : ℕ} (f : Vec a → Vec b) :

Prop

f scales with its argument: f (s • v) = s • f v, spelled pointwise. The statement of every cotangent-chain _smul lemma; an abbrev, so rw [h] and h s v see the equation.

Equations

Proofs.IsHomog f = ∀ (s : ℝ) (v : Proofs.Vec a), (f fun (i : Fin a) => s * v i) = fun (i : Fin b) => s * f v i

Instances For

source

theorem Proofs.IsHomog.comp {a b c : ℕ} {g : Vec b → Vec c} {f : Vec a → Vec b} (hg : IsHomog g) (hf : IsHomog f) :

IsHomog (g ∘ f)

source

theorem Proofs.HasVJP.backward_smul {m n : ℕ} {f : Vec m → Vec n} (hf : HasVJP f) (x : Vec m) :

IsHomog (hf.backward x)

A VJP backward is linear in its cotangent — read off HasVJP.correct: it scales.

source

theorem Proofs.HasVJP.backward_add {m n : ℕ} {f : Vec m → Vec n} (hf : HasVJP f) (x : Vec m) (dy dy' : Vec n) :

hf.backward x (dy + dy') = hf.backward x dy + hf.backward x dy'

A VJP backward is linear in its cotangent — it adds.

source

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

HasVJP fun (y : Vec a) (k : Fin b) => y (σ k)

The VJP of a coordinate reindex y ↦ y ∘ σ (a gather, reindexCLM σ): the backward scatters each output cotangent back to the input cell it was read from — pdiv_reindex's indicator, contracted. The one witness behind every reindex-shaped layer: the BN layout bridges (reassocFwd/Back, bnchwFwd/Back), the StridedConv decimations, and reindexHasVJP's correct. (broadcastFlatHasVJP keeps its own witness: its backward is spelled the way the IR's broadcastBack denotes, and the SE-gate graph ties match it by rfl.)

Equations

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

Instances For

source

theorem Proofs.reindexVJP_backward_of_inv {a b : ℕ} (σ : Fin b → Fin a) (τ : Fin a → Fin b) (hστ : ∀ (i : Fin a), σ (τ i) = i) (hτσ : ∀ (k : Fin b), τ (σ k) = k) (v : Vec a) (dy : Vec b) :

(reindexVJP σ).backward v dy = fun (i : Fin a) => dy (τ i)

Along a bijection the scatter is the inverse gather: when τ inverts σ, exactly one delta survives, so reindexVJP σ's backward is the reindex along τ.

source

theorem Proofs.reindexVJP_backward {a b : ℕ} (σ : Fin b → Fin a) (v : Vec a) (dy : Vec b) :

(reindexVJP σ).backward v dy = fun (i : Fin a) => ∑ k : Fin b, if i = σ k then dy k else 0

The scatter written as a masked sum (if i = σ k then dy k else 0) — the spelling the IR's scatter ops denote.

source

@[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

Instances For

source

noncomputable def Proofs.biPathHasVJP {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.biPathHasVJP 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 := ⋯ }

Instances For

source

@[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

Instances For

source

noncomputable def Proofs.elemwiseProductHasVJP {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.

Instances For

source

def Proofs.identityHasVJP (n : ℕ) :

HasVJP fun (x : Vec n) => x

Identity VJP — proved, no sorry.

Equations

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

Instances For

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.

source

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

Type

A vector–Jacobian product for f at one point x: a backward map with the proof that it contracts f's Jacobian (pdiv f x) against the cotangent. The pointwise form of HasVJP, for operators that are differentiable only away from their kinks.

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

Instances For

source

theorem Proofs.HasVJPAt.ext_iff {m n : ℕ} {f : Vec m → Vec n} {x : Vec m} {x✝ y : HasVJPAt f x} :

x✝ = y ↔ x✝.backward = y.backward

source

theorem Proofs.HasVJPAt.ext {m n : ℕ} {f : Vec m → Vec n} {x : Vec m} {x✝ y : HasVJPAt f x} (backward : x✝.backward = y.backward) :

x✝ = y

source

instance Proofs.instSubsingletonHasVJPAt {m n : ℕ} {f : Vec m → Vec n} {x : Vec m} :

Subsingleton (HasVJPAt f x)

A map has at most one VJP witness at a point.

source

noncomputable def Proofs.HasVJPAt.canonical {m n : ℕ} (f : Vec m → Vec n) (x : Vec m) :

HasVJPAt f x

The canonical witness at a point: the backward is the pdiv contraction itself.

Equations

Proofs.HasVJPAt.canonical f x = { backward := fun (dy : Proofs.Vec n) (i : Fin m) => ∑ j : Fin n, Proofs.pdiv f x i j * dy j, correct := ⋯ }

Instances For

source

def Proofs.HasVJPAt.congr {m n : ℕ} {f g : Vec m → Vec n} {x : Vec m} (h : f = g) (hf : HasVJPAt f x) :

HasVJPAt g x

Transport a pointwise witness along an equality of maps; the backward is carried over verbatim (HasVJP.congr's pointwise peer).

Equations

Proofs.HasVJPAt.congr h hf = { backward := hf.backward, correct := ⋯ }

Instances For

source

@[simp]

theorem Proofs.HasVJPAt.congr_backward {m n : ℕ} {f g : Vec m → Vec n} {x : Vec m} (h : f = g) (hf : HasVJPAt f x) :

(congr h hf).backward = hf.backward

source

theorem Proofs.HasVJPAt.backward_unique_of_eq {m n : ℕ} {f g : Vec m → Vec n} {x : Vec m} (hfg : f = g) (h₁ : HasVJPAt f x) (h₂ : HasVJPAt g x) (dy : Vec n) :

h₁.backward dy = h₂.backward dy

Two HasVJPAt witnesses for EQUAL maps at one point have the same backward — the pointwise peer of HasVJP.backward_unique_of_eq, through .correct rather than a transport.

source

theorem Proofs.HasVJPAt.backward_unique {m n : ℕ} {f : Vec m → Vec n} {x : Vec m} (h₁ h₂ : HasVJPAt f x) (dy : Vec n) :

h₁.backward dy = h₂.backward dy

Any two HasVJPAt witnesses for the same map at the same point have the same backward — HasVJP.backward_unique's pointwise peer.

source

theorem Proofs.HasVJPAt.backward_smul {m n : ℕ} {f : Vec m → Vec n} {x : Vec m} (hf : HasVJPAt f x) :

IsHomog hf.backward

A pointwise VJP backward scales with its cotangent (HasVJP.backward_smul's peer).

source

theorem Proofs.HasVJPAt.backward_add {m n : ℕ} {f : Vec m → Vec n} {x : Vec m} (hf : HasVJPAt f x) (dy dy' : Vec n) :

hf.backward (dy + dy') = hf.backward dy + hf.backward dy'

A pointwise VJP backward adds (HasVJP.backward_add's peer).

source

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 := ⋯ }

Instances For

source

def Proofs.identityHasVJPAt (n : ℕ) (x : Vec n) :

HasVJPAt (fun (y : Vec n) => y) x

Identity pointwise VJP — trivial.

Equations

Proofs.identityHasVJPAt n x = (Proofs.identityHasVJP n).toHasVJPAt x

Instances For

source

noncomputable def Proofs.vjpCompAt {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 vjpComp, 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.vjpCompAt f g x hf_diff hg_diff hf hg = { backward := fun (dy : Proofs.Vec p) => hf.backward (hg.backward dy), correct := ⋯ }

Instances For

source

noncomputable def Proofs.vjpComp {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 — vjpCompAt at every point. Requires f and g to be differentiable everywhere.

Equations

Proofs.vjpComp 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 := ⋯ }

Instances For

source

theorem Proofs.vjpCompAt_backward {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)) (dy : Vec p) :

(vjpCompAt f g x hf_diff hg_diff hf hg).backward dy = hf.backward (hg.backward dy)

vjpCompAt's backward: g's, then f's. Definitional; stated so a chain peels by rw. Note: Not by simp only in a whole-net tie: simp uses an rfl lemma as a dsimp step and the kernel re-derives the unfolding.

source

theorem Proofs.vjpComp_backward {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) (x : Vec m) (dy : Vec p) :

(vjpComp f g hf_diff hg_diff hf hg).backward x dy = hf.backward x (hg.backward (f x) dy)

vjpComp's backward: g's at f x, then f's. Definitional, for rw.

source

@[reducible, inline]

abbrev Proofs.HasVJPDiffAt {m n : ℕ} (f : Vec m → Vec n) (x : Vec m) :

Type

A VJP witness at x together with differentiability there — one stage of a whole-net chain. A PProd, so the DifferentiableAt Prop can ride along with the data.

Equations

Proofs.HasVJPDiffAt f x = (Proofs.HasVJPAt f x ×' DifferentiableAt ℝ f x)

Instances For

source

noncomputable def Proofs.vjpCompDiffAt {m n p : ℕ} (f : Vec m → Vec n) (g : Vec n → Vec p) (x : Vec m) (hf : HasVJPDiffAt f x) (hg : HasVJPDiffAt g (f x)) :

HasVJPDiffAt (g ∘ f) x

Compose two HasVJPDiffAt stages. The fold step of every whole-net apex.

Equations

Proofs.vjpCompDiffAt f g x hf hg = ⟨Proofs.vjpCompAt f g x ⋯ ⋯ hf.fst hg.fst, ⋯⟩

Instances For

source

theorem Proofs.vjpCompDiffAt_fst_backward {m n p : ℕ} (f : Vec m → Vec n) (g : Vec n → Vec p) (x : Vec m) (hf : HasVJPDiffAt f x) (hg : HasVJPDiffAt g (f x)) (dy : Vec p) :

(vjpCompDiffAt f g x hf hg).fst.backward dy = hf.fst.backward (hg.fst.backward dy)

One vjpCompDiffAt level's backward, unfolded: the composite runs g's backward, then f's. Definitional, stated so that a chain of levels peels by rw rather than by a rfl that has to find the same unfolding through concrete witnesses. Note: Not by simp only: simp would use it as a dsimp step and record nothing for the kernel to replay.

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.

source

theorem Proofs.sum_finProdFinEquiv {M : Type u_1} [AddCommMonoid M] {m n : ℕ} (f : Fin (m * n) → M) :

∑ k : Fin (m * n), f k = ∑ i : Fin m, ∑ j : Fin n, f (finProdFinEquiv (i, j))

Row-major reindexing of a flat sum — a sum over Fin (m * n) is the double sum over (i, j) read through finProdFinEquiv. Mathlib has the Fin m × Fin n form (Fintype.sum_prod_type) but no Fin (m * n) one; every flatten/unflatten sum in the suite reduces to this split.

source

theorem Proofs.sum_finProdFinEquiv₃ {M : Type u_1} [AddCommMonoid M] {a b c : ℕ} (f : Fin (a * b * c) → M) :

∑ k : Fin (a * b * c), f k = ∑ i : Fin a, ∑ j : Fin b, ∑ l : Fin c, f (finProdFinEquiv (finProdFinEquiv (i, j), l))

The rank-3 case: a Fin (a·b·c) sum is the row-major triple sum — the layout Tensor3.flatten and every c·h·w activation index use.

source

theorem Proofs.sum_channel_fiber (oc m : ℕ) (idx : Fin oc) (g : Fin (oc * m) → ℝ) :

(∑ j : Fin (oc * m), if idx = (finProdFinEquiv.symm j).1 then g j else 0) = ∑ s : Fin m, g (finProdFinEquiv (idx, s))

Sum-over-the-channel-fibre: Σ_j [idx = chan j]·g j = Σ_s g (idx, s).

source

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

Instances For

source

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

Instances For

source

theorem Proofs.Mat.flatten_apply {m n : ℕ} (A : Mat m n) (k : Fin (m * n)) :

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

source

theorem Proofs.Mat.unflatten_apply {m n : ℕ} (v : Vec (m * n)) (i : Fin m) (j : Fin n) :

unflatten v i j = v (finProdFinEquiv (i, j))

source

@[simp]

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

unflatten A.flatten = A

Unflatten is a left inverse of flatten.

source

@[simp]

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

(unflatten v).flatten = v

Flatten is a left inverse of unflatten.

source

@[simp]

theorem Proofs.Mat.flatten_finProdFinEquiv {m n : ℕ} (A : Mat m n) (i : Fin m) (j : Fin n) :

A.flatten (finProdFinEquiv (i, j)) = A i j

Reading a flattened matrix at an encoded index reads the entry.

source

theorem Proofs.Mat.flatten_differentiable {m n : ℕ} :

Differentiable ℝ flatten

Mat.flatten is a coordinate reindexing, hence differentiable.

source

theorem Proofs.Mat.unflatten_differentiable {m n : ℕ} :

Differentiable ℝ unflatten

Mat.unflatten is a coordinate reindexing, hence differentiable.

source

theorem Proofs.Mat.transpose_differentiable {m n : ℕ} :

Differentiable ℝ transpose

Mat.transpose is a coordinate reindexing, hence differentiable.

source

theorem Proofs.Mat.mul_const_differentiable {m p q : ℕ} (D : Mat p q) :

Differentiable ℝ fun (A : Mat m p) => A.mul D

Right-multiplying by a constant matrix is differentiable.

source

theorem Proofs.Mat.const_mul_differentiable {m p q : ℕ} (C : Mat m p) :

Differentiable ℝ fun (B : Mat p q) => C.mul B

Left-multiplying by a constant matrix is differentiable.

source

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. The rank-2 structural rules (chain/sum/id) 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))

Instances For

source

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 — derived from pdiv_comp via the row-major flatten bijection.

source

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.

source

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.

source

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

Instances For

source

theorem Proofs.HasVJPMat.ext_iff {a b c d : ℕ} {f : Mat a b → Mat c d} {x y : HasVJPMat f} :

x = y ↔ x.backward = y.backward

source

theorem Proofs.HasVJPMat.ext {a b c d : ℕ} {f : Mat a b → Mat c d} {x y : HasVJPMat f} (backward : x.backward = y.backward) :

x = y

source

instance Proofs.instSubsingletonHasVJPMat {a b c d : ℕ} {f : Mat a b → Mat c d} :

Subsingleton (HasVJPMat f)

A matrix map has at most one VJP witness.

source

noncomputable def Proofs.HasVJPMat.canonical {a b c d : ℕ} (f : Mat a b → Mat c d) :

HasVJPMat f

The canonical matrix witness: the backward is the pdivMat contraction itself.

Equations

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

Instances For

source

theorem Proofs.HasVJPMat.backward_unique_of_eq {a b c d : ℕ} {f g : Mat a b → Mat c d} (hfg : f = g) (v : HasVJPMat f) (v' : HasVJPMat g) (A : Mat a b) (dY : Mat c d) :

v.backward A dY = v'.backward A dY

Two HasVJPMat witnesses for EQUAL maps have the same backward — the matrix peer of HasVJP.backward_unique_of_eq, through .correct rather than a transport.

source

noncomputable def Proofs.vjpMatComp {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 vjpComp to rank-2 indices.

Equations

Proofs.vjpMatComp 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 := ⋯ }

Instances For

source

@[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

Instances For

source

noncomputable def Proofs.biPathMatHasVJP {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.

Instances For

source

noncomputable def Proofs.identityMatHasVJP (a b : ℕ) :

HasVJPMat fun (M : Mat a b) => M

Identity VJP for matrices — proved, no sorry.

Equations

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

Instances For

source

noncomputable def Proofs.HasVJPMat.toHasVJP {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 (vitBody, 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.

Instances For

source

theorem Proofs.HasVJPMat.toHasVJP_backward {a b c d : ℕ} {f : Mat a b → Mat c d} (hf : HasVJPMat f) (v : Vec (a * b)) (dy : Vec (c * d)) (idx : Fin (a * b)) :

hf.toHasVJP.backward v dy idx = hf.backward (Mat.unflatten v) (Mat.unflatten dy) (finProdFinEquiv.symm idx).1 (finProdFinEquiv.symm idx).2

HasVJPMat.toHasVJP's backward: decode the flat index, run the matrix backward.

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

pdivMat_matmul_left_const — right-factor varies, left factor fixed: ∂(C · B')_{kl} / ∂B'_{ij} = C_{ki} · [l = j].
pdivMat_matmul_right_const — left factor varies, right factor fixed: ∂(A' · D)_{kl} / ∂A'_{ij} = D_{jl} · [i = k].
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.

source

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) — B' ↦ C·B' is linear, so by pdiv_of_linear the entry is C·E_{ij} read at (k, l): C k i when l = j.

source

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) — the left-const case with roles swapped: A' ↦ A'·D is linear, and E_{ij}·D read at (k, l) is D j l when i = k.

source

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

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

Block-diagonal Jacobian of a per-row family, at a point. Applying g r to each row r keeps the matrix Jacobian block-diagonal across rows: the (i, j, k, l) entry is pdiv (g k) (A k) j l when i = k and 0 otherwise. Each g r need only be differentiable at its own row A r: the flat map's coordinate (r, l) is g r's coordinate l after the row projection, and hasFDerivAt_pi assembles the rows' derivatives.

source

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.

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: pdivMat_rowIndep_perRow_at with one g for every row.

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.

source

noncomputable def Proofs.rowwiseHasVJPMat {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 (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 rowSoftmaxHasVJPMat: any per-token operation (LayerNorm, GELU, dense, activation) lifts to a per-sequence matrix operation via this one helper.

Equations

Proofs.rowwiseHasVJPMat 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 := ⋯ }

Instances For

source

theorem Proofs.rowwiseHasVJPMat_backward {m n p : ℕ} {g : Vec n → Vec p} (hg : HasVJP g) (hg_diff : Differentiable ℝ g) (A : Mat m n) (dY : Mat m p) (r : Fin m) :

(rowwiseHasVJPMat hg hg_diff).backward A dY r = hg.backward (A r) (dY r)

rowwiseHasVJPMat's backward: the row map's backward, row by row.

source

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_const_smul + pdiv_id via the flatten bijection. ∂(s · A')_{kl} / ∂A'_{ij} = s · δ_{ik,jl}.

source

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

source

noncomputable def Proofs.matmulLeftConstHasVJP {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.matmulLeftConstHasVJP 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 := ⋯ }

Instances For

source

noncomputable def Proofs.matmulRightConstHasVJP {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.matmulRightConstHasVJP 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 := ⋯ }

Instances For

source

noncomputable def Proofs.scalarScaleHasVJP {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.scalarScaleHasVJP s = { backward := fun (_A dY : Proofs.Mat m n) (i : Fin m) (j : Fin n) => s * dY i j, correct := ⋯ }

Instances For

source

noncomputable def Proofs.transposeHasVJP {m n : ℕ} :

HasVJPMat fun (M : Mat m n) => M.transpose

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

Equations

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

Instances For

source

@[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 → ℝ)

Instances For

source

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

Instances For

source

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

Instances For

source

theorem Proofs.Tensor3.flatten_apply {c h w : ℕ} (T : Tensor3 c h w) (k : Fin (c * h * w)) :

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

source

theorem Proofs.Tensor3.unflatten_apply {c h w : ℕ} (v : Vec (c * h * w)) (ci : Fin c) (hi : Fin h) (wi : Fin w) :

unflatten v ci hi wi = v (finProdFinEquiv (finProdFinEquiv (ci, hi), wi))

source

@[simp]

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

unflatten T.flatten = T

source

@[simp]

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

(unflatten v).flatten = v

source

@[simp]

theorem Proofs.Tensor3.flatten_finProdFinEquiv {c h w : ℕ} (T : Tensor3 c h w) (ci : Fin c) (hi : Fin h) (wi : Fin w) :

T.flatten (finProdFinEquiv (finProdFinEquiv (ci, hi), wi)) = T ci hi wi

Reading a flattened tensor at an encoded index reads the entry.

source

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.

source

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)), hence a projection.

source

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 — pdiv of the flattened map, both indices read through the triple-nested flatten bijection. Operator-specific VJPs at rank 3 (conv2dHasVJP3, maxPool2HasVJP3, depthwiseHasVJP3) are bundled HasVJP3 defs in their own files; nets compose them at Vec, through HasVJP3.toHasVJP / HasVJPAt3.toHasVJPAt.

Equations

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

Instances For

source

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

Instances For

source

theorem Proofs.HasVJP3.ext_iff {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} {f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂} {x y : HasVJP3 f} :

x = y ↔ x.backward = y.backward

source

theorem Proofs.HasVJP3.ext {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} {f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂} {x y : HasVJP3 f} (backward : x.backward = y.backward) :

x = y

source

instance Proofs.instSubsingletonHasVJP3 {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} {f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂} :

Subsingleton (HasVJP3 f)

A rank-3 map has at most one VJP witness.

source

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

HasVJP3 f

The canonical rank-3 witness: the backward is the pdiv3 contraction itself.

Equations

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

Instances For

source

theorem Proofs.HasVJP3.backward_smul {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} {f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂} (hf : HasVJP3 f) (x : Tensor3 c₁ h₁ w₁) (a : ℝ) (dy : Tensor3 c₂ h₂ w₂) :

(hf.backward x fun (i₁ : Fin c₂) (i₂ : Fin h₂) (i₃ : Fin w₂) => a * dy i₁ i₂ i₃) = fun (j₁ : Fin c₁) (j₂ : Fin h₁) (j₃ : Fin w₁) => a * hf.backward x dy j₁ j₂ j₃

A rank-3 VJP backward scales with its cotangent (HasVJP.backward_smul's peer).

source

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 maxPool2HasVJPAt3 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

Instances For

source

theorem Proofs.HasVJPAt3.ext_iff {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} {f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂} {x : Tensor3 c₁ h₁ w₁} {x✝ y : HasVJPAt3 f x} :

x✝ = y ↔ x✝.backward = y.backward

source

theorem Proofs.HasVJPAt3.ext {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} {f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂} {x : Tensor3 c₁ h₁ w₁} {x✝ y : HasVJPAt3 f x} (backward : x✝.backward = y.backward) :

x✝ = y

source

instance Proofs.instSubsingletonHasVJPAt3 {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} {f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂} {x : Tensor3 c₁ h₁ w₁} :

Subsingleton (HasVJPAt3 f x)

A rank-3 map has at most one VJP witness at a point.

source

noncomputable def Proofs.HasVJPAt3.canonical {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} (f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂) (x : Tensor3 c₁ h₁ w₁) :

HasVJPAt3 f x

The canonical rank-3 witness at a point: the backward is the pdiv3 contraction itself.

Equations

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

Instances For

source

noncomputable def Proofs.HasVJP3.toHasVJP {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.toHasVJP. 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 sum_finProdFinEquiv twice.

Equations

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

Instances For

source

theorem Proofs.HasVJP3.toHasVJP_backward {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} {f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂} (hf : HasVJP3 f) (v : Vec (c₁ * h₁ * w₁)) (dy : Vec (c₂ * h₂ * w₂)) (idx : Fin (c₁ * h₁ * w₁)) :

hf.toHasVJP.backward v dy idx = (hf.backward (Tensor3.unflatten v) (Tensor3.unflatten dy)).flatten idx

HasVJP3.toHasVJP's backward: decode the flat index, run the rank-3 backward.

source

theorem Proofs.HasVJP3.flatten_backward {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} {f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂} (hf : HasVJP3 f) (x : Tensor3 c₁ h₁ w₁) (dy : Vec (c₂ * h₂ * w₂)) (idx : Fin (c₁ * h₁ * w₁)) :

(hf.backward x (Tensor3.unflatten dy)).flatten idx = ∑ j : Fin (c₂ * h₂ * w₂), pdiv (fun (v : Vec (c₁ * h₁ * w₁)) => (f (Tensor3.unflatten v)).flatten) x.flatten idx j * dy j

A rank-3 backward, flattened, is the flat pdiv contraction — HasVJP3.toHasVJP's correct, read at a tensor activation.

source

noncomputable def Proofs.HasVJPAt3.toHasVJPAt {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.toHasVJP, 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.

Instances For

source

theorem Proofs.HasVJPAt3.toHasVJPAt_backward {c₁ h₁ w₁ c₂ h₂ w₂ : ℕ} {f : Tensor3 c₁ h₁ w₁ → Tensor3 c₂ h₂ w₂} {x : Tensor3 c₁ h₁ w₁} (hf : HasVJPAt3 f x) (dy : Vec (c₂ * h₂ * w₂)) (idx : Fin (c₁ * h₁ * w₁)) :

hf.toHasVJPAt.backward dy idx = (hf.backward (Tensor3.unflatten dy)).flatten idx

HasVJPAt3.toHasVJPAt's backward: decode the flat index, run the rank-3 backward.

Documentation

LeanMlir.Proofs.Foundation.Tensor

Tensor algebra and the VJP framework — the root of the proof suite #