Pullbacks of higher model derivatives

The public reverse transform seeds input directions while holding model state constant, runs the checked reverse pass, and returns state and input gradients. The proof connects that complete IO call to the adjoint derivative of a higher input derivative. It uses the existing flattened context only at the calculus boundary; callers still receive the usual typed state and tensor.

Successful recording and execution are explicit hypotheses, separate from the first-order graph certificate, higher-order jet laws, and smoothness needed to interchange derivatives. No theorem here identifies rounded floating-point execution with exact real differentiation.

theorem TorchLean.autograd.model.vjp_eq_checked {σ τ : Shape} (model : nn.Sequential σ τ) {α : Type} [Storage α] [Context α] (state : State model α) (input : Tensor α σ) (seed : Tensor α τ) (graph : nn.TypedGraphModel (Runtime.Autograd.Model.Layers.Seq.stateShapes model) σ τ α) (valid : Runtime.Autograd.Model.Layers.Seq.validate model = Except.ok ()) (lowered : Runtime.Autograd.Model.Layers.Seq.lowerToTypedGraph model Runtime.Autograd.Model.Layers.Mode.eval = pure graph) (result : TensorPack α (Runtime.Autograd.Model.Layers.Seq.stateShapes model ++ [σ]) × Tensor α τ) (checked : Runtime.Autograd.Torch.TypedGraphWithData.vjpChecked graph ((nn.State.Internal.toTensorPack state).append (TensorPack.singleton input)) () seed = Except.ok result) : vjp model state input seed = pure (nn.State.Internal.fromTensorPack result.1.split.1, result.1.split.2.head)

Validation and a successful checked graph execution determine the public VJP result.

This execution equation works with every runtime scalar context. It preserves the existing state/input split and makes no claim that the graph's backward rules are analytic derivatives.

theorem TorchLean.autograd.model.derivativeVjp_eq {σ τ : Shape} (model : nn.Sequential σ τ) (state : State model ℝ) (input : Tensor ℝ σ) (directions : List (Tensor ℝ σ)) (seed : Tensor ℝ τ) (real : nn.TypedGraphModel (Runtime.Autograd.Model.Layers.Seq.stateShapes model) σ τ ℝ) (nested : nn.TypedGraphModel (Runtime.Autograd.Model.Layers.Seq.stateShapes model) σ τ (Runtime.Autograd.Model.Dual.Nested ℝ directions.length)) (valid : Runtime.Autograd.Model.Layers.Seq.validate model = Except.ok ()) (lowered : Runtime.Autograd.Model.Layers.Seq.lowerToTypedGraph model Runtime.Autograd.Model.Layers.Mode.eval = pure nested) (sameShapes : nested.nodeShapes = real.nodeShapes) (hgraph : Proofs.Autograd.Algebra.GraphData.PreservesPullbackJet (Proofs.Autograd.CtxVec (Runtime.Autograd.Model.Layers.Seq.stateShapes model ++ [σ])) directions.length real.data (sameShapes ▸ nested.data)) (sameOutput : sameShapes ▸ nested.output = real.output) (proofGraph : Proofs.Autograd.Algebra.Graph Unit (Runtime.Autograd.Model.Layers.Seq.stateShapes model ++ [σ]) real.nodeShapes) (same : proofGraph.toData = real.data) (correct : (y : Proofs.Autograd.CtxVec (Runtime.Autograd.Model.Layers.Seq.stateShapes model ++ [σ])) → Proofs.Autograd.GraphFDerivCorrectAt (proofGraph.toReal ()) y) (smooth : ContDiff ℝ (↑directions.length + 1) fun (z : Proofs.Autograd.CtxVec (Runtime.Autograd.Model.Layers.Seq.stateShapes model ++ [σ])) => Proofs.Autograd.tensorToVec (Runtime.Autograd.Torch.TypedGraphWithData.forward real (Proofs.Autograd.unflattenCtx z) ())) (result : TensorPack (Runtime.Autograd.Model.Dual.Nested ℝ directions.length) (Runtime.Autograd.Model.Layers.Seq.stateShapes model ++ [σ]) × Tensor (Runtime.Autograd.Model.Dual.Nested ℝ directions.length) τ) (checked : Runtime.Autograd.Torch.TypedGraphWithData.vjpChecked nested ((TensorPack.map (fun {shape : Shape} => Tensor.map (Runtime.Autograd.Model.Dual.Nested.ofPrimal directions.length)) (nn.State.Internal.toTensorPack state)).append (TensorPack.singleton (Runtime.Autograd.Model.Dual.Nested.seedTensor (fun (i : Fin directions.length) => directions[i]) input))) () (Tensor.map (Runtime.Autograd.Model.Dual.Nested.ofPrimal directions.length) seed) = Except.ok result) : have ds := fun (i : Fin directions.length) => Proofs.Autograd.flattenCtx (TensorPack.zero.append (TensorPack.singleton directions[i])); have x := Proofs.Autograd.flattenCtx ((nn.State.Internal.toTensorPack state).append (TensorPack.singleton input)); have gradient := Proofs.Autograd.unflattenCtx ((ContinuousLinearMap.adjoint (fderiv ℝ (fun (y : Proofs.Autograd.CtxVec (Runtime.Autograd.Model.Layers.Seq.stateShapes model ++ [σ])) => (iteratedFDeriv ℝ directions.length (fun (z : Proofs.Autograd.CtxVec (Runtime.Autograd.Model.Layers.Seq.stateShapes model ++ [σ])) => Proofs.Autograd.tensorToVec (Runtime.Autograd.Torch.TypedGraphWithData.forward real (Proofs.Autograd.unflattenCtx z) ())) y) ds) x)) (Proofs.Autograd.tensorToVec seed)); derivativeVjp model state input directions seed = pure (nn.State.Internal.fromTensorPack gradient.split.1, gradient.split.2.head)

The public higher-order VJP returns the pullback of the requested iterated input derivative.

The context includes state and input entries. Directions have zero state components, but the final adjoint differentiates with respect to all entries, giving both parameter and input gradients. The output cotangent and directions are fixed, including when directions repeat.