Tape-node context primitives

This module contains the low-level vectorized context operations used by the tape-node proof library: block projections, one-hot cotangent injections, and the bridge from generic OpSpecFDerivCorrect witnesses to NodeFDerivCorrect nodes.

Coordinate lemmas for the Euclidean identification

Mathlib's EuclideanSpace is PiLp 2, a type synonym carrying a WithLp wrapper. The four lemmas below are the plumbing that lets us forget the wrapper: each says that reading coordinate i of a vector built from f gives f i, for the various shapes the wrapper takes. They are boring on purpose, and having them as simp lemmas is what keeps the real proofs in this file readable.

@[simp]

theorem Proofs.Autograd.piLpContinuousLinearEquiv2_symm_apply {n : ℕ} (f : Fin n → ℝ) (i : Fin n) : ((PiLp.continuousLinearEquiv 2 ℝ fun (x : Fin n) => ℝ).symm f).ofLp i = f i

Coordinates of the inverse PiLp equivalence are the values of the underlying function.

@[simp]

theorem Proofs.Autograd.piLpContinuousLinearEquiv2_symm_clm_apply {n : ℕ} (f : Fin n → ℝ) (i : Fin n) : (↑(PiLp.continuousLinearEquiv 2 ℝ fun (x : Fin n) => ℝ).symm f).ofLp i = f i

The same, applied through the bundled continuous linear map.

@[simp]

theorem Proofs.Autograd.piLpContinuousLinearEquiv2_symm_clm_apply_ofLp {n : ℕ} (f : Fin n → ℝ) (i : Fin n) : (↑(PiLp.continuousLinearEquiv 2 ℝ fun (x : Fin n) => ℝ).symm f).ofLp i = f i

The same again, with the ofLp projection made explicit.

@[simp]

theorem Proofs.Autograd.euclideanEquiv_symm_ofLp {n : ℕ} (f : Fin n → ℝ) (i : Fin n) : ((EuclideanSpace.equiv (Fin n) ℝ).symm f).ofLp i = f i

And once more for EuclideanSpace.equiv, the spelling used by vecOfFun.

@[simp]

theorem Proofs.Autograd.inner_scalarVec_left (a : ℝ) (δ : Vec Spec.Shape.scalar.size) : inner ℝ (vecOfFun fun (x : Fin Spec.Shape.scalar.size) => a) δ = a * δ.ofLp ⟨0, inner_scalarVec_left._proof_1⟩

Inner product against a one-dimensional constant vector collapses to a single product.

Scalar tensors vectorize to Vec 1, so this is the lemma that turns the adjointness statement for a scalar node into ordinary multiplication instead of a sum over Finset.univ.

def Proofs.Autograd.CtxVec.getBlock {Γ : List Spec.Shape} (i : Fin Γ.length) : CtxVec Γ → Vec (Γ.get i).size

Project a vectorized context onto the block at list position i.

CtxVec.get is the shape-indexed interface; it accepts an Idx Γ s and transports the result to the statically known shape s.

def Proofs.Autograd.CtxVec.singleBlock {Γ : List Spec.Shape} (i : Fin Γ.length) : Vec (Γ.get i).size → CtxVec Γ

Inject v into block i of a vectorized context and fill every other block with zeros.

This is the adjoint of getBlock with respect to the Euclidean inner product.

theorem Proofs.Autograd.CtxVec.inner_getBlock_singleBlock {Γ : List Spec.Shape} (i : Fin Γ.length) (x : CtxVec Γ) (v : Vec (Γ.get i).size) : inner ℝ x (singleBlock i v) = inner ℝ (getBlock i x) v

Adjointness of block projection and injection.

This is the vectorized counterpart of the one-hot cotangent principle used in tape soundness.

noncomputable def Proofs.Autograd.CtxVec.get {Γ : List Spec.Shape} {s : Spec.Shape} (idx : Idx Γ s) (x : CtxVec Γ) : Vec s.size

Project the block specified by idx : Idx Γ s out of a vectorized context.

theorem Proofs.Autograd.CtxVec.getBlock_flattenCtx {Γ : List Spec.Shape} (xs : TorchLean.TensorPack ℝ Γ) (i : Fin Γ.length) : getBlock i (flattenCtx xs) = tensorToVec (xs.get i)

Selecting a block after vectorization agrees with selecting the original tensor.

theorem Proofs.Autograd.CtxVec.get_flattenCtx {Γ : List Spec.Shape} {s : Spec.Shape} (idx : Idx Γ s) (xs : TorchLean.TensorPack ℝ Γ) : get idx (flattenCtx xs) = tensorToVec (getIdx xs idx)

Typed tensor selection commutes with vectorization, including its shape cast.

noncomputable def Proofs.Autograd.CtxVec.single {Γ : List Spec.Shape} {s : Spec.Shape} (idx : Idx Γ s) (v : Vec s.size) : CtxVec Γ

Inject a block into a vectorized context at idx, filling other blocks with zeros.

theorem Proofs.Autograd.CtxVec.inner_get_single {Γ : List Spec.Shape} {s : Spec.Shape} (idx : Idx Γ s) (x : CtxVec Γ) (v : Vec s.size) : inner ℝ x (single idx v) = inner ℝ (get idx x) v

Adjointness of get/single: ⟪x, single idx v⟫ = ⟪get idx x, v⟫.

noncomputable def Proofs.Autograd.CtxVec.headCLM {s : Spec.Shape} {ss : List Spec.Shape} : CtxVec (s :: ss) →L[ℝ] Vec s.size

Continuous linear map extracting the head block of a nonempty vectorized context.

@[simp]

theorem Proofs.Autograd.CtxVec.headCLM_apply {s : Spec.Shape} {ss : List Spec.Shape} (x : CtxVec (s :: ss)) (j : Fin s.size) : (headCLM x).ofLp j = x.ofLp (Fin.castAdd (ctxSize ss) j)

headCLM reads the leading block of coordinates, exactly as flattenCtx_cons lays them out.

noncomputable def Proofs.Autograd.CtxVec.tailCLM {s : Spec.Shape} {ss : List Spec.Shape} : CtxVec (s :: ss) →L[ℝ] CtxVec ss

Continuous linear map extracting the tail blocks of a nonempty vectorized context.

@[simp]

theorem Proofs.Autograd.CtxVec.tailCLM_apply {s : Spec.Shape} {ss : List Spec.Shape} (x : CtxVec (s :: ss)) (j : Fin (ctxSize ss)) : (tailCLM x).ofLp j = x.ofLp (Fin.natAdd s.size j)

tailCLM reads the trailing blocks.

def Proofs.Autograd.CtxVec.getBlockCLM {Γ : List Spec.Shape} (i : Fin Γ.length) : CtxVec Γ →L[ℝ] Vec (Γ.get i).size

getBlock as a continuous linear map, constructed recursively from headCLM and tailCLM.

@[simp]

theorem Proofs.Autograd.CtxVec.getBlockCLM_apply {Γ : List Spec.Shape} (i : Fin Γ.length) (x : CtxVec Γ) : (getBlockCLM i) x = getBlock i x

The bundled getBlockCLM computes the same thing as the plain getBlock.

This is the payoff of building it recursively from headCLM and tailCLM: block selection comes out continuous and linear by construction, so nothing downstream has to prove it again.

noncomputable def Proofs.Autograd.CtxVec.getCLM {Γ : List Spec.Shape} {s : Spec.Shape} (idx : Idx Γ s) : CtxVec Γ →L[ℝ] Vec s.size

get packaged as a continuous linear map.

@[simp]

theorem Proofs.Autograd.CtxVec.getCLM_apply {Γ : List Spec.Shape} {s : Spec.Shape} (idx : Idx Γ s) (x : CtxVec Γ) : (getCLM idx) x = get idx x

And the shape-indexed getCLM agrees with get.

Nodes in this file are authored directly on the vectorized context CtxVec.

This is the most convenient authoring style for analytic proofs: forwardVec/jvpVec/vjpVec are definitional, and the correctness obligation is an inner-product identity on Euclidean vectors.

noncomputable def Proofs.Autograd.Node.ofFn {Γ : List Spec.Shape} {τ : Spec.Shape} (f : CtxVec Γ → Vec τ.size) (jvp : CtxVec Γ → CtxVec Γ → Vec τ.size) (vjp : CtxVec Γ → Vec τ.size → CtxVec Γ) (correct_inner : ∀ (x dx : CtxVec Γ) (δ : Vec τ.size), inner ℝ (jvp x dx) δ = inner ℝ dx (vjp x δ)) : Node Γ τ

Convenience constructor: build a tape Node from vector-level forward/JVP/VJP plus adjointness.

The correct_inner field is exactly the local VJP/JVP law: ⟪jvp x dx, δ⟫ = ⟪dx, vjp x δ⟫.

@[simp]

theorem Proofs.Autograd.Node.forwardVec_ofFn {Γ : List Spec.Shape} {τ : Spec.Shape} (f : CtxVec Γ → Vec τ.size) (jvp : CtxVec Γ → CtxVec Γ → Vec τ.size) (vjp : CtxVec Γ → Vec τ.size → CtxVec Γ) (h : ∀ (x dx : CtxVec Γ) (δ : Vec τ.size), inner ℝ (jvp x dx) δ = inner ℝ dx (vjp x δ)) : (ofFn f jvp vjp h).forwardVec = f

A node built by ofFn has the given forward map.

@[simp]

theorem Proofs.Autograd.Node.jvpVec_ofFn {Γ : List Spec.Shape} {τ : Spec.Shape} (f : CtxVec Γ → Vec τ.size) (jvp : CtxVec Γ → CtxVec Γ → Vec τ.size) (vjp : CtxVec Γ → Vec τ.size → CtxVec Γ) (h : ∀ (x dx : CtxVec Γ) (δ : Vec τ.size), inner ℝ (jvp x dx) δ = inner ℝ dx (vjp x δ)) : (ofFn f jvp vjp h).jvpVec = jvp

A node built by ofFn has the given forward-mode derivative.

@[simp]

theorem Proofs.Autograd.Node.vjpVec_ofFn {Γ : List Spec.Shape} {τ : Spec.Shape} (f : CtxVec Γ → Vec τ.size) (jvp : CtxVec Γ → CtxVec Γ → Vec τ.size) (vjp : CtxVec Γ → Vec τ.size → CtxVec Γ) (h : ∀ (x dx : CtxVec Γ) (δ : Vec τ.size), inner ℝ (jvp x dx) δ = inner ℝ dx (vjp x δ)) : (ofFn f jvp vjp h).vjpVec = vjp

A node built by ofFn has the given reverse-mode derivative. Together the three projection lemmas mean a caller never has to unfold ofFn, only supply the soundness argument h once.

noncomputable def Proofs.Autograd.OpSpecFDerivCorrect.linear {inDim outDim : ℕ} (m : Spec.LinearSpec ℝ inDim outDim) : OpSpecFDerivCorrect inDim outDim

OpSpecFDerivCorrect instance for a linear layer.

This is the analytic correctness lemma behind the tape node constructors: it identifies the JVP with the Fréchet derivative (a matrix multiplication) for linearSpec.

PyTorch analogue: the torch.nn.Linear forward map is affine, so its derivative is constant. https://pytorch.org/docs/stable/generated/torch.nn.Linear.html