FDeriv Core

HasFDerivAt-level (analytic) soundness for the proved-correct autograd layer.

This file starts by connecting our tensor dot to the Euclidean-space inner product, sets up the generic coordinatewise calculus on Euclidean vectors (elemwiseVec, elemwiseDerivCLM), then proves a first end-to-end theorem for a 2-layer MLP (Linear → ReLU → Linear):

• the OpSpec reverse-mode backward computes the true analytic VJP, i.e. backward x δ = VJP[f, x] δ (after translating between tensors and vectors).

Notes:

• Everything here is over ℝ (spec-level exact arithmetic).
• ReLU is not differentiable at 0, so the theorems assume a "no kinks" hypothesis on the pre-activation vector.
• The tensor-output theorem shape is naturally VJP-based: for f : ℝⁿ → ℝᵐ, reverse-mode computes δ ↦ (Df(x))ᵗ δ. Scalar losses are the special case m = 1 / δ = 1.

PyTorch correspondence / citations

• Reverse-mode VJPs and Jacobian-transpose products are exactly what PyTorch’s backward computes. https://pytorch.org/docs/stable/autograd.html
• Linear layers and ReLU as used in the example are standard PyTorch building blocks. https://pytorch.org/docs/stable/generated/torch.nn.Linear.html https://pytorch.org/docs/stable/generated/torch.nn.functional.relu.html

@[reducible, inline]

noncomputable abbrev Proofs.Autograd.euclideanEquiv (n : ℕ) : EuclideanSpace ℝ (Fin n) ≃L[ℝ] Fin n → ℝ

Abbreviation for the Euclidean-space equivalence Vec n ≃ (Fin n → ℝ).

Dot-product vs. Euclidean inner product

To connect OpSpecCorrect (stated with the tensor dot product) to fderiv and adjoints (stated with Euclidean inner products), we prove that Spec.dot agrees with inner after vectorization.

@[simp]

theorem Proofs.Autograd.euclideanEquiv_getScalarE {n : ℕ} (t : TorchLean.Tensor ℝ [n]) : (euclideanEquiv n) (getScalarE t) = t.getScalar

getScalarE is defined via EuclideanSpace.equiv; this lemma exposes the underlying coordinates.

theorem Proofs.Autograd.dot_eq_inner_vec {n : ℕ} (a b : TorchLean.Tensor ℝ [n]) : Spec.dot a b = inner ℝ (getScalarE a) (getScalarE b)

For 1D scalar tensors, Spec.dot agrees with the Euclidean inner product on Vec n after converting via getScalarE.

theorem Proofs.Autograd.getScalarE_add_spec {n : ℕ} (a b : TorchLean.Tensor ℝ [n]) : getScalarE (a.addSpec b) = getScalarE a + getScalarE b

Vectorization commutes with tensor addition.

theorem Proofs.Autograd.getScalarE_map_spec {n : ℕ} (f : ℝ → ℝ) (t : TorchLean.Tensor ℝ [n]) : getScalarE (TorchLean.Tensor.map f t) = (euclideanEquiv n).symm fun (i : Fin n) => f (t.getScalar i)

Vectorization commutes with elementwise mapping: getScalarE (Tensor.map f t) applies f to each coordinate of TorchLean.Tensor.getScalar t.

theorem Proofs.Autograd.getScalarE_relu_deriv_spec {n : ℕ} (t : TorchLean.Tensor ℝ [n]) : getScalarE (Activation.reluDerivSpec t) = (euclideanEquiv n).symm fun (i : Fin n) => Activation.Math.reluDerivSpec (t.getScalar i)

Vectorization of reluDerivSpec: the derivative mask is ReLU’s scalar derivative applied coordinatewise.

def Proofs.Autograd.tensorToMatrix {m n : ℕ} (W : TorchLean.Tensor ℝ [m, n]) : Matrix (Fin m) (Fin n) ℝ

View a matrix-shaped tensor W : Tensor ℝ (m×n) as a Mathlib Matrix (Fin m) (Fin n) ℝ.

This is just the coordinate function Spec.get2.

noncomputable def Proofs.Autograd.matCLM {m n : ℕ} (W : Matrix (Fin m) (Fin n) ℝ) : Vec n →L[ℝ] Vec m

The matrix–vector multiplication map as a continuous linear map on Euclidean vectors.

This is the Euclidean-space version of the tensor op matVecMulSpec.

theorem Proofs.Autograd.getScalarE_mat_vec_mul_spec {m n : ℕ} (A : TorchLean.Tensor ℝ [m, n]) (v : TorchLean.Tensor ℝ [n]) : getScalarE (Spec.matVecMulSpec A v) = (matCLM (tensorToMatrix A)) (getScalarE v)

Vectorization commutes with matrix–vector multiplication: getScalarE (mat_vec_mul_spec A v) = (matCLM (tensorToMatrix A)) (getScalarE v).

noncomputable def Proofs.Autograd.affine {inDim outDim : ℕ} (W : Matrix (Fin outDim) (Fin inDim) ℝ) (b : Vec outDim) : Vec inDim → Vec outDim

Affine map x ↦ W x + b on Euclidean vectors.

This is the vector-space analogue of Spec.linearSpec.

theorem Proofs.Autograd.hasFDerivAt_affine {inDim outDim : ℕ} (W : Matrix (Fin outDim) (Fin inDim) ℝ) (b : Vec outDim) (x : Vec inDim) : HasFDerivAt (affine W b) (matCLM W) x

affine is Fréchet-differentiable with derivative W (as a CLM), since it is linear + constant.

theorem Proofs.Autograd.getScalarE_linear_spec {inDim outDim : ℕ} (l : Spec.LinearSpec ℝ inDim outDim) (x : TorchLean.Tensor ℝ [inDim]) : getScalarE (Spec.linearSpec l x) = affine (tensorToMatrix l.weights) (getScalarE l.bias) (getScalarE x)

Vectorization of Spec.linearSpec is the Euclidean affine map built from the same weights/bias.

def Proofs.Autograd.elemwiseVec {n : ℕ} (f : ℝ → ℝ) : Vec n → Vec n

Apply a scalar function f : ℝ → ℝ coordinatewise to a vector.

This is the Euclidean-space analogue of the tensor-level Tensor.map.

noncomputable def Proofs.Autograd.evalCLM {n : ℕ} (i : Fin n) : Vec n →L[ℝ] ℝ

Coordinate evaluation as a continuous linear map on Vec n.

@[simp]

theorem Proofs.Autograd.evalCLM_apply {n : ℕ} (i : Fin n) (x : Vec n) : (evalCLM i) x = x.ofLp i

The coordinate evaluation functional reads coordinate i.

noncomputable def Proofs.Autograd.elemwiseDerivCLM {n : ℕ} (f' : ℝ → ℝ) (x : Vec n) : Vec n →L[ℝ] Vec n

The derivative candidate for elemwiseVec f at a point x, built from a proposed scalar derivative f'.

Concretely: (elemwiseDerivCLM f' x) dx has coordinates i ↦ f'(xᵢ) * dxᵢ.

theorem Proofs.Autograd.hasFDerivAt_elemwiseVec_at {n : ℕ} {f f' : ℝ → ℝ} (x : Vec n) (hf : ∀ (i : Fin n), HasDerivAt f (f' (x.ofLp i)) (x.ofLp i)) : HasFDerivAt (elemwiseVec f) (elemwiseDerivCLM f' x) x

elemwiseVec f is Fréchet differentiable at x as soon as f is differentiable at every coordinate of x, with the diagonal derivative elemwiseDerivCLM f' x.

This is the form needed for maps such as ReLU whose scalar derivative exists only away from finitely many points.

theorem Proofs.Autograd.hasFDerivAt_elemwiseVec {n : ℕ} {f f' : ℝ → ℝ} (x : Vec n) (hf : ∀ (z : ℝ), HasDerivAt f (f' z) z) : HasFDerivAt (elemwiseVec f) (elemwiseDerivCLM f' x) x

If f is differentiable everywhere with derivative f', then elemwiseVec f is Fréchet differentiable everywhere with derivative elemwiseDerivCLM f'.

noncomputable def Proofs.Autograd.reluVec {n : ℕ} : Vec n → Vec n

ReLU as a map on Euclidean vectors (coordinatewise max x 0).

This is the Euclidean-space analogue of Spec.relu_op.forward.

noncomputable def Proofs.Autograd.reluDerivCLM {n : ℕ} (x : Vec n) : Vec n →L[ℝ] Vec n

Derivative of reluVec at x: the diagonal scaling by the scalar ReLU derivative mask.

theorem Proofs.Autograd.hasFDerivAt_reluVec {n : ℕ} (x : Vec n) (hx : ∀ (i : Fin n), x.ofLp i ≠ 0) : HasFDerivAt reluVec (reluDerivCLM x) x

ReLU is not differentiable at 0, so its Fréchet derivative is asserted under the “no kinks” hypothesis that every coordinate of x is nonzero.

noncomputable def Proofs.Autograd.mlpVec {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec ℝ inDim hidDim) (l2 : Spec.LinearSpec ℝ hidDim outDim) : Vec inDim → Vec outDim

2-layer MLP forward map on Euclidean vectors:

x ↦ affine W2 b2 (relu (affine W1 b1 x)).

noncomputable def Proofs.Autograd.mlpDeriv {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec ℝ inDim hidDim) (l2 : Spec.LinearSpec ℝ hidDim outDim) (x : Vec inDim) : Vec inDim →L[ℝ] Vec outDim

Closed-form derivative (as a continuous linear map) of mlpVec at x.

This is the chain rule composition: W2 ∘ ReLU'(z1) ∘ W1.

theorem Proofs.Autograd.hasFDerivAt_mlpVec {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec ℝ inDim hidDim) (l2 : Spec.LinearSpec ℝ hidDim outDim) (x : Vec inDim) (hx : ∀ (i : Fin hidDim), have W1 := tensorToMatrix l1.weights; have b1 := getScalarE l1.bias; (affine W1 b1 x).ofLp i ≠ 0) : HasFDerivAt (mlpVec l1 l2) (mlpDeriv l1 l2 x) x

Fréchet differentiability of the 2-layer MLP (Linear → ReLU → Linear) under a “no kinks” hypothesis.

Because ReLU is not differentiable at 0, we assume all pre-activation coordinates z1ᵢ are nonzero.

noncomputable def Proofs.Autograd.mlpOp {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec ℝ inDim hidDim) (l2 : Spec.LinearSpec ℝ hidDim outDim) : Spec.OpSpec ℝ (Spec.Shape.dim inDim Spec.Shape.scalar) (Spec.Shape.dim outDim Spec.Shape.scalar)

The spec-level MLP as a composed Spec.OpSpec:

linear l1 then relu then linear l2.

noncomputable def Proofs.Autograd.mlpCorrect {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec ℝ inDim hidDim) (l2 : Spec.LinearSpec ℝ hidDim outDim) : OpSpecCorrect (Spec.Shape.dim inDim Spec.Shape.scalar) (Spec.Shape.dim outDim Spec.Shape.scalar)

The proved-correct MLP OpSpecCorrect, built by composing the primitive correctness lemmas.

This provides the dot-level adjointness statement: ⟪JVP,δ⟫ = ⟪dx,VJP⟫.

theorem Proofs.Autograd.getScalar_mlp_jvp {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec ℝ inDim hidDim) (l2 : Spec.LinearSpec ℝ hidDim outDim) (x dx : TorchLean.Tensor ℝ [inDim]) : getScalarE ((mlpCorrect l1 l2).jvp x dx) = (mlpDeriv l1 l2 (getScalarE x)) (getScalarE dx)

Identify the OpSpecCorrect JVP for the MLP with the analytic derivative mlpDeriv, after vectorizing tensors to Euclidean vectors.

theorem Proofs.Autograd.mlp_backward_eq_adjoint_fderiv {inDim hidDim outDim : ℕ} (l1 : Spec.LinearSpec ℝ inDim hidDim) (l2 : Spec.LinearSpec ℝ hidDim outDim) (x : TorchLean.Tensor ℝ [inDim]) (hx : ∀ (i : Fin hidDim), have W1 := tensorToMatrix l1.weights; have b1 := getScalarE l1.bias; (affine W1 b1 (getScalarE x)).ofLp i ≠ 0) (δ : TorchLean.Tensor ℝ [outDim]) : getScalarE ((mlpOp l1 l2).backward x δ) = (vjp (mlpVec l1 l2) (getScalarE x)) (getScalarE δ)

End-to-end analytic soundness for the 2-layer MLP OpSpec:

the OpSpec.backward returned by the spec-level reverse-mode rule equals the adjoint of the true Fréchet derivative of the forward map (i.e. the analytic VJP), after vectorization.

This is the proof layer analogue of PyTorch’s claim that loss.backward() computes the correct VJP for the composed model, assuming the primitive backward rules are correct.