Higher derivatives through reverse accumulation

A reverse pass reads saved activations, computes local pullbacks, and adds their contributions. The jet relation tracks all three operations, with arbitrary tensor shapes and finite derivative order. Seeds can depend smoothly on the inputs, as they do when differentiating a composed loss.

These theorems use the existing GraphData.backpropCtx. They identify nested-dual execution with derivatives of its real implementation. Identifying that implementation with the adjoint of the forward derivative requires the separate first-order correctness certificate.

theorem TorchLean.TensorPack.JetRelated.zero {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {shapes : List Spec.Shape} (directions : Fin n → E) (x : E) : JetRelated directions x (fun (x : E) => TensorPack.zero) TensorPack.zero

Zero-initialized gradient buffers contain no nonzero derivative coefficients.

theorem TorchLean.TensorPack.JetRelated.single {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {shapes : List Spec.Shape} {shape : Spec.Shape} (index : Proofs.Idx shapes shape) (directions : Fin n → E) (x : E) {f : E → Tensor ℝ shape} (hf : ContDiff ℝ (↑n) f) : JetRelated directions x (fun (y : E) => Proofs.Autograd.Algebra.TensorPack.single index (f y)) (Proofs.Autograd.Algebra.TensorPack.single index (Runtime.Autograd.Model.DualTensor.jet directions f x))

Place a tensor's full jet at one typed input and zero-fill the other gradient buffers.

theorem TorchLean.TensorPack.JetRelated.add {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {shapes : List Spec.Shape} {directions : Fin n → E} {x : E} {f g : E → TensorPack ℝ shapes} {fv gv : TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) shapes} (hf : JetRelated directions x f fv) (hg : JetRelated directions x g gv) : JetRelated directions x (fun (y : E) => (f y).add (g y)) (fv.add gv)

Gradient accumulation preserves mixed derivatives, including shared-parent contributions.

theorem TorchLean.TensorPack.JetRelated.unsnoc {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {shapes : List Spec.Shape} {shape : Spec.Shape} {directions : Fin n → E} {x : E} {f : E → TensorPack ℝ (shapes ++ [shape])} {values : TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) (shapes ++ [shape])} (h : JetRelated directions x f values) : JetRelated directions x (fun (y : E) => (f y).unsnoc.1) values.unsnoc.1 ∧ (ContDiff ℝ ↑n fun (y : E) => (f y).unsnoc.2) ∧ values.unsnoc.2 = Runtime.Autograd.Model.DualTensor.jet directions (fun (y : E) => (f y).unsnoc.2) x

Splitting off a node's cotangent preserves both parts of the derivative information.

def Proofs.Autograd.Algebra.NodeData.PreservesPullbackJet (E : Type u_1) [NormedAddCommGroup E] [NormedSpace ℝ E] (n : ℕ) {Δ : Type} {Γ : List Spec.Shape} {shape : Spec.Shape} (real : TorchLean.TensorPack ℝ Γ → Δ → TorchLean.Tensor ℝ shape → TorchLean.TensorPack ℝ Γ) (nested : TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) Γ → Δ → TorchLean.Tensor (Runtime.Autograd.Model.Dual.Nested ℝ n) shape → TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) Γ) : Prop

The implemented pullback preserves full jets of activations and smoothly varying seeds. This law concerns differentiation of the pullback; its first-order adjoint law is separate.

theorem Proofs.Autograd.Algebra.NodeData.PreservesPullbackJet.single {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {Δ : Type} {Γ : List Spec.Shape} {shape : Spec.Shape} (index : Idx Γ shape) : PreservesPullbackJet E n (fun (x : TorchLean.TensorPack ℝ Γ) (x_1 : Δ) (seed : TorchLean.Tensor ℝ shape) => TensorPack.single index seed) fun (x : TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) Γ) (x_1 : Δ) (seed : TorchLean.Tensor (Runtime.Autograd.Model.Dual.Nested ℝ n) shape) => TensorPack.single index seed

Passing a seed directly to one input preserves its full jet.

theorem Proofs.Autograd.Algebra.NodeData.PreservesPullbackJet.add {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {Δ : Type} {Γ : List Spec.Shape} {shape : Spec.Shape} {realLeft realRight : TorchLean.TensorPack ℝ Γ → Δ → TorchLean.Tensor ℝ shape → TorchLean.TensorPack ℝ Γ} {nestedLeft nestedRight : TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) Γ → Δ → TorchLean.Tensor (Runtime.Autograd.Model.Dual.Nested ℝ n) shape → TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) Γ} (hleft : PreservesPullbackJet E n realLeft nestedLeft) (hright : PreservesPullbackJet E n realRight nestedRight) : PreservesPullbackJet E n (fun (ctx : TorchLean.TensorPack ℝ Γ) (data : Δ) (seed : TorchLean.Tensor ℝ shape) => (realLeft ctx data seed).add (realRight ctx data seed)) fun (ctx : TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) Γ) (data : Δ) (seed : TorchLean.Tensor (Runtime.Autograd.Model.Dual.Nested ℝ n) shape) => (nestedLeft ctx data seed).add (nestedRight ctx data seed)

Sum independently certified contributions, even when they target the same input.

theorem Proofs.Autograd.Algebra.NodeData.PreservesPullbackJet.mul_left {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {Δ : Type} {Γ : List Spec.Shape} {shape : Spec.Shape} (index : Idx Γ shape) {real : TorchLean.TensorPack ℝ Γ → Δ → TorchLean.Tensor ℝ shape} {nested : TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) Γ → Δ → TorchLean.Tensor (Runtime.Autograd.Model.Dual.Nested ℝ n) shape} (h : PreservesJet E n real nested) : PreservesPullbackJet E n (fun (ctx : TorchLean.TensorPack ℝ Γ) (data : Δ) (seed : TorchLean.Tensor ℝ shape) => TensorPack.single index (seed.mulSpec (real ctx data))) fun (ctx : TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) Γ) (data : Δ) (seed : TorchLean.Tensor (Runtime.Autograd.Model.Dual.Nested ℝ n) shape) => TensorPack.single index (seed.mulSpec (nested ctx data))

Scatter seed * coefficient to one typed input, preserving all mixed derivatives.

theorem Proofs.Autograd.Algebra.NodeData.PreservesPullbackJet.mul_right {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {Δ : Type} {Γ : List Spec.Shape} {shape : Spec.Shape} (index : Idx Γ shape) {real : TorchLean.TensorPack ℝ Γ → Δ → TorchLean.Tensor ℝ shape} {nested : TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) Γ → Δ → TorchLean.Tensor (Runtime.Autograd.Model.Dual.Nested ℝ n) shape} (h : PreservesJet E n real nested) : PreservesPullbackJet E n (fun (ctx : TorchLean.TensorPack ℝ Γ) (data : Δ) (seed : TorchLean.Tensor ℝ shape) => TensorPack.single index ((real ctx data).mulSpec seed)) fun (ctx : TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) Γ) (data : Δ) (seed : TorchLean.Tensor (Runtime.Autograd.Model.Dual.Nested ℝ n) shape) => TensorPack.single index ((nested ctx data).mulSpec seed)

Scatter coefficient * seed to one typed input, preserving all mixed derivatives.

inductive Proofs.Autograd.Algebra.GraphData.PreservesPullbackJet (E : Type u_1) [NormedAddCommGroup E] [NormedSpace ℝ E] (n : ℕ) {Δ : Type} {Γ shapes : List Spec.Shape} : GraphData ℝ Δ Γ shapes → GraphData (Runtime.Autograd.Model.Dual.Nested ℝ n) Δ Γ shapes → Prop

Local jet laws for both saved forward values and the pullbacks in an executable graph.

• nil {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {Δ : Type} {Γ : List Spec.Shape} : PreservesPullbackJet E n GraphData.nil GraphData.nil

  The empty reverse pass returns its supplied cotangents.

• snoc {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {Δ : Type} {Γ shapes : List Spec.Shape} {shape : Spec.Shape} {real : GraphData ℝ Δ Γ shapes} {nested : GraphData (Runtime.Autograd.Model.Dual.Nested ℝ n) Δ Γ shapes} {realNode : NodeData ℝ Δ (Γ ++ shapes) shape} {nestedNode : NodeData (Runtime.Autograd.Model.Dual.Nested ℝ n) Δ (Γ ++ shapes) shape} (previous : PreservesPullbackJet E n real nested) (forward : NodeData.PreservesJet E n realNode.forward nestedNode.forward) (pullback : NodeData.PreservesPullbackJet E n realNode.vjp nestedNode.vjp) : PreservesPullbackJet E n (real.snoc realNode) (nested.snoc nestedNode)

  Extend a graph with independent forward and pullback jet laws.

theorem Proofs.Autograd.Algebra.GraphData.PreservesPullbackJet.forward {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {Δ : Type} {Γ shapes : List Spec.Shape} {real : GraphData ℝ Δ Γ shapes} {nested : GraphData (Runtime.Autograd.Model.Dual.Nested ℝ n) Δ Γ shapes} (h : PreservesPullbackJet E n real nested) : PreservesJet E n real nested

Reverse-mode certificates also certify the saved activations used by each pullback.

theorem Proofs.Autograd.Algebra.GraphData.PreservesPullbackJet.backprop {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {Δ : Type} {Γ shapes : List Spec.Shape} {real : GraphData ℝ Δ Γ shapes} {nested : GraphData (Runtime.Autograd.Model.Dual.Nested ℝ n) Δ Γ shapes} (h : PreservesPullbackJet E n real nested) (directions : Fin n → E) (x : E) (inputs : E → TorchLean.TensorPack ℝ Γ) (values : TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) Γ) (data : Δ) (hinputs : TorchLean.TensorPack.JetRelated directions x inputs values) (seed : E → TorchLean.TensorPack ℝ (Γ ++ shapes)) (seedValues : TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) (Γ ++ shapes)) (hseed : TorchLean.TensorPack.JetRelated directions x seed seedValues) : TorchLean.TensorPack.JetRelated directions x (fun (y : E) => real.backpropCtx (inputs y) data (seed y)) (nested.backpropCtx values data seedValues)

Reverse accumulation preserves all derivative coefficients of every input cotangent.

theorem Proofs.Autograd.Algebra.GraphData.PreservesPullbackJet.tangent_backprop {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {Δ : Type} {Γ shapes : List Spec.Shape} {real : GraphData ℝ Δ Γ shapes} {nested : GraphData (Runtime.Autograd.Model.Dual.Nested ℝ n) Δ Γ shapes} (h : PreservesPullbackJet E n real nested) (directions : Fin n → E) (x : E) (inputs : E → TorchLean.TensorPack ℝ Γ) (values : TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) Γ) (data : Δ) (hinputs : TorchLean.TensorPack.JetRelated directions x inputs values) (seed : E → TorchLean.TensorPack ℝ (Γ ++ shapes)) (seedValues : TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) (Γ ++ shapes)) (hseed : TorchLean.TensorPack.JetRelated directions x seed seedValues) {shape : Spec.Shape} (input : Idx Γ shape) : Runtime.Autograd.Model.Dual.Nested.tangentTensor (getIdx (nested.backpropCtx values data seedValues) input) = (iteratedFDeriv ℝ n (fun (y : E) => getIdx (real.backpropCtx (inputs y) data (seed y)) input) x) directions

Extracting a reverse result gives the iterated derivative of the implemented real pullback. The seed may vary with the input; a fixed loss cotangent is a special case.