Coordinatewise primitive specifications

A captured tensor can supply a different scalar function at every coordinate. Its derivative is still diagonal: each input tangent is multiplied by the corresponding scalar derivative. Over the reals this diagonal map is self-adjoint, which gives the reverse rule for the same scalar formula.

The final lemmas apply to an actual OpSpec. They require its forward and backward coordinate formulas, as well as scalar derivative evidence at the input being differentiated. This keeps the analytic step separate from an algebraic pairing identity and permits the necessary domain conditions for reciprocals and piecewise functions. Flattening uses the existing tensor/vector equivalence, so the statements apply to arbitrary tensor shapes, including empty ones.

noncomputable def Proofs.Autograd.PrimitiveSpecs.coordinatewise {n : ℕ} (f : Fin n → ℝ → ℝ) (x : Vec n) : Vec n

Apply a possibly different scalar function at each coordinate.

noncomputable def Proofs.Autograd.PrimitiveSpecs.coordinateDeriv {n : ℕ} (coefficients : Fin n → ℝ) : Vec n →L[ℝ] Vec n

The diagonal linear map with the supplied scalar coefficients.

@[simp]

theorem Proofs.Autograd.PrimitiveSpecs.coordinateDeriv_apply {n : ℕ} (coefficients : Fin n → ℝ) (dx : Vec n) (i : Fin n) : ((coordinateDeriv coefficients) dx).ofLp i = coefficients i * dx.ofLp i

A diagonal derivative multiplies each tangent by its own coefficient.

theorem Proofs.Autograd.PrimitiveSpecs.hasFDerivAt_coordinatewise {n : ℕ} (f : Fin n → ℝ → ℝ) (coefficients : Fin n → ℝ) (x : Vec n) (hf : ∀ (i : Fin n), HasDerivAt (f i) (coefficients i) (x.ofLp i)) : HasFDerivAt (coordinatewise f) (coordinateDeriv coefficients) x

Scalar derivatives at the input coordinates determine the derivative of the whole vector.

theorem Proofs.Autograd.PrimitiveSpecs.coordinateDeriv_adjoint_apply {n : ℕ} (coefficients : Fin n → ℝ) (δ : Vec n) : (ContinuousLinearMap.adjoint (coordinateDeriv coefficients)) δ = (coordinateDeriv coefficients) δ

Real diagonal maps are self-adjoint, independently of differentiability.

theorem Proofs.Autograd.PrimitiveSpecs.hasFDerivAt_of_coordinates {s : Spec.Shape} (op : Spec.OpSpec ℝ s s) (f : Fin s.size → ℝ → ℝ) (coefficients : Fin s.size → ℝ) (x : Vec s.size) (hforward : ∀ (y : Vec s.size), tensorToVec (op.forward (vecToTensor y)) = coordinatewise f y) (hf : ∀ (i : Fin s.size), HasDerivAt (f i) (coefficients i) (x.ofLp i)) : HasFDerivAt (fun (y : Vec s.size) => tensorToVec (op.forward (vecToTensor y))) (coordinateDeriv coefficients) x

Transport coordinatewise calculus to the exact forward function stored by an OpSpec.

theorem Proofs.Autograd.PrimitiveSpecs.backward_eq_adjoint_of_coordinates {s : Spec.Shape} (op : Spec.OpSpec ℝ s s) (f : Fin s.size → ℝ → ℝ) (coefficients : Fin s.size → ℝ) (x : Vec s.size) (hforward : ∀ (y : Vec s.size), tensorToVec (op.forward (vecToTensor y)) = coordinatewise f y) (hf : ∀ (i : Fin s.size), HasDerivAt (f i) (coefficients i) (x.ofLp i)) (hbackward : ∀ (δ : TorchLean.Tensor ℝ s), tensorToVec (op.backward (vecToTensor x) δ) = (coordinateDeriv coefficients) (tensorToVec δ)) (δ : TorchLean.Tensor ℝ s) : tensorToVec (op.backward (vecToTensor x) δ) = (ContinuousLinearMap.adjoint (fderiv ℝ (fun (y : Vec s.size) => tensorToVec (op.forward (vecToTensor y))) x)) (tensorToVec δ)

Identify the stored backward with the adjoint of the actual forward derivative.

The backward formula is checked at the same input as the scalar derivative hypotheses. In particular, supplying a selected slope at a kink does not establish a classical derivative there.