Higher-order graph evaluation

The real and nested-dual executions use the existing GraphData representation. A local jet-preservation law for each pair of nodes propagates through the graph, including its saved intermediates. Selecting any output then gives mathlib's iterated Fréchet derivative.

The auxiliary environment is held fixed. These results concern exact-real forward evaluation; they do not assume that node JVP/VJP fields are correct merely because the forward maps agree.

def TorchLean.TensorPack.JetRelated {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} (directions : Fin n → E) (x : E) {shapes : List Spec.Shape} : (E → TensorPack ℝ shapes) → TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) shapes → Prop

Each saved tensor is smooth and its nested-dual value contains its mixed derivatives.

theorem TorchLean.TensorPack.JetRelated.const {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {shapes : List Spec.Shape} (directions : Fin n → E) (x : E) (state : TensorPack ℝ shapes) : JetRelated directions x (fun (x : E) => state) (map (fun {shape : Spec.Shape} => Tensor.map (Runtime.Autograd.Model.Dual.Nested.ofPrimal n)) state)

Fixed model state has zero mixed derivatives and uses the runtime constant embedding.

theorem TorchLean.TensorPack.JetRelated.singleton_seed {n : ℕ} {shape : Spec.Shape} (directions : Fin n → Tensor ℝ shape) (input : Tensor ℝ shape) : JetRelated directions input (fun (x : Tensor ℝ shape) => singleton x) (singleton (Runtime.Autograd.Model.Dual.Nested.seedTensor directions input))

The model's seeded input tensor satisfies the context relation at every derivative order.

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

Join separately certified input families, for example fixed state and varying model inputs.

theorem TorchLean.TensorPack.JetRelated.cast {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {shapes target : List Spec.Shape} {directions : Fin n → E} {x : E} {f : E → TensorPack ℝ shapes} {values : TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) shapes} (h : JetRelated directions x f values) (hs : shapes = target) : JetRelated directions x (fun (y : E) => TensorPack.cast hs (f y)) (TensorPack.cast hs values)

Shape-list transport changes neither smoothness nor stored derivative coefficients.

theorem TorchLean.TensorPack.JetRelated.snoc {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} {values : TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) shapes} (h : JetRelated directions x f values) {g : E → Tensor ℝ shape} (hg : ContDiff ℝ (↑n) g) : JetRelated directions x (fun (y : E) => (f y).snoc (g y)) (values.snoc (Runtime.Autograd.Model.DualTensor.jet directions g x))

Appending a proved node output extends the relation for the saved context.

theorem TorchLean.TensorPack.JetRelated.get {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} {values : TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) shapes} (h : JetRelated directions x f values) (index : Proofs.Idx shapes shape) : (ContDiff ℝ ↑n fun (y : E) => Proofs.getIdx (f y) index) ∧ Proofs.getIdx values index = Runtime.Autograd.Model.DualTensor.jet directions (fun (y : E) => Proofs.getIdx (f y) index) x

A typed context lookup retrieves a smooth tensor and its complete jet.

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

Coordinatewise certificates determine the complete heterogeneous context relation.

def Proofs.Autograd.Algebra.NodeData.PreservesJet (E : Type u_1) [NormedAddCommGroup E] [NormedSpace ℝ E] (n : ℕ) {Δ : Type} {Γ : List Spec.Shape} {shape : Spec.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) : Prop

A pair of forward maps preserves smooth input families and their full mixed-derivative jets. Reverse-mode correctness is a separate obligation.

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

Reading an existing context entry preserves its smoothness and all derivative coefficients.

theorem Proofs.Autograd.Algebra.NodeData.PreservesJet.map {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {Δ : Type} {Γ : List Spec.Shape} {shape : Spec.Shape} (index : Idx Γ shape) (f : ℝ → ℝ) (lifted : Runtime.Autograd.Model.Dual.Nested ℝ n → Runtime.Autograd.Model.Dual.Nested ℝ n) (hf : ContDiff ℝ (↑n) f) (hjet : ∀ (directions : Fin n → E) (g : E → ℝ), ContDiff ℝ (↑n) g → ∀ (x : E), Runtime.Autograd.Model.Dual.jet directions (fun (y : E) => f (g y)) x = lifted (Runtime.Autograd.Model.Dual.jet directions g x)) : PreservesJet E n (fun (inputs : TorchLean.TensorPack ℝ Γ) (x : Δ) => TorchLean.Tensor.map f (getIdx inputs index)) fun (values : TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) Γ) (x : Δ) => TorchLean.Tensor.map lifted (getIdx values index)

Lift a scalar jet rule through an elementwise graph operation reading a saved tensor.

theorem Proofs.Autograd.Algebra.NodeData.PreservesJet.map2 {E : Type u_1} [NormedAddCommGroup E] [NormedSpace ℝ E] {n : ℕ} {Δ : Type} {Γ : List Spec.Shape} {shape : Spec.Shape} (left right : Idx Γ shape) (f : ℝ → ℝ → ℝ) (lifted : Runtime.Autograd.Model.Dual.Nested ℝ n → Runtime.Autograd.Model.Dual.Nested ℝ n → Runtime.Autograd.Model.Dual.Nested ℝ n) (hf : ContDiff ℝ ↑n fun (p : ℝ × ℝ) => f p.1 p.2) (hjet : ∀ (directions : Fin n → E) (g h : E → ℝ), ContDiff ℝ (↑n) g → ContDiff ℝ (↑n) h → ∀ (x : E), Runtime.Autograd.Model.Dual.jet directions (fun (y : E) => f (g y) (h y)) x = lifted (Runtime.Autograd.Model.Dual.jet directions g x) (Runtime.Autograd.Model.Dual.jet directions h x)) : PreservesJet E n (fun (inputs : TorchLean.TensorPack ℝ Γ) (x : Δ) => TorchLean.Tensor.map2Spec f (getIdx inputs left) (getIdx inputs right)) fun (values : TorchLean.TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ n) Γ) (x : Δ) => TorchLean.Tensor.map2Spec lifted (getIdx values left) (getIdx values right)

Lift a binary scalar jet rule through any pair of saved tensors of the same shape.

inductive Proofs.Autograd.Algebra.GraphData.PreservesJet (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

Node-local jet laws on a pair of existing executable graphs, in their recorded order.

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

  An empty graph preserves the supplied input jets.

• 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 : PreservesJet E n real nested) (node : NodeData.PreservesJet E n realNode.forward nestedNode.forward) : PreservesJet E n (real.snoc realNode) (nested.snoc nestedNode)

  Appending a locally proved node preserves all existing and newly computed jets.

theorem Proofs.Autograd.Algebra.GraphData.PreservesJet.eval {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 : PreservesJet 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) : TorchLean.TensorPack.JetRelated directions x (fun (y : E) => real.eval (inputs y) data) (nested.eval values data)

Graph evaluation preserves the jet relation for the entire saved context.

theorem Proofs.Autograd.Algebra.GraphData.PreservesJet.tangent_eval {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 : PreservesJet 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) {shape : Spec.Shape} (output : Idx (Γ ++ shapes) shape) : Runtime.Autograd.Model.Dual.Nested.tangentTensor (getIdx (nested.eval values data) output) = (iteratedFDeriv ℝ n (fun (y : E) => getIdx (real.eval (inputs y) data) output) x) directions

Extracting all tangent coefficients at any graph output computes its iterated derivative.