Autograd Transforms

Advanced differentiation operations are kept out of the first autograd tutorial:

• jacrev and hessian for tensor functions;
• jvp and hvp for model state;
• Loss.detach for an explicit gradient stop;
• repeated and mixed input derivatives of a smooth model;
• composing a custom operation's proof with checked first-order differentiation;
• arbitrary-order derivatives of a recorded graph;
• a complete model derivative call, including successful lowering.

Build this module after NN.Examples.Quickstart.AutogradBasics.

Run:

scripts/lake.sh exe torchlean autograd_transforms

def NN.Examples.DeepDives.AutogradTransforms.exeName : String

Subcommand name, used in the usage text and in argument-error messages.

def NN.Examples.DeepDives.AutogradTransforms.square {shape : Spec.Shape} : TorchLean.autograd.Function shape shape

Componentwise squaring, whose Jacobian is the diagonal matrix 2x. Small enough that the printed rows can be checked by hand.

def NN.Examples.DeepDives.AutogradTransforms.meanSquare {shape : Spec.Shape} : TorchLean.autograd.Function shape []

Mean of the squares: a scalar-valued function, so it has a Hessian to compute.

def NN.Examples.DeepDives.AutogradTransforms.model : TorchLean.nn.Sequential [2] [3]

One linear layer 2 -> 3, seeded deterministically so the printed numbers are reproducible.

def NN.Examples.DeepDives.AutogradTransforms.runDemo : IO Unit

Run higher-order and directional differentiation examples.

def NN.Examples.DeepDives.AutogradTransforms.usage : String

Help text; the demo takes no flags.

def NN.Examples.DeepDives.AutogradTransforms.main (args : List String) : IO Unit

Entry point.

noncomputable def NN.Examples.DeepDives.AutogradTransforms.Correctness.square {Γ ss : List Spec.Shape} {s : Spec.Shape} (graph : Proofs.Autograd.DGraph Γ ss) (input : Proofs.Idx (Γ ++ ss) s) : Proofs.Autograd.DGraph Γ (ss ++ [s])

Add elementwise squaring to a proved graph, for any tensor shape.

autograd composes mathlib's scalar rules and lifts them to tensors; snoc carries the certificate into the composed graph. No shape-specific proof is needed.

noncomputable def NN.Examples.DeepDives.AutogradTransforms.Correctness.fourthPower (s : Spec.Shape) : Proofs.Autograd.DGraph [s] [s, s]

Squaring twice reuses the same primitive proof; no new backward rule is needed.

def NN.Examples.DeepDives.AutogradTransforms.Correctness.output (s : Spec.Shape) : Proofs.Idx ([s] ++ [s, s]) s

The second square is the output; the input and first square remain saved graph values.

theorem NN.Examples.DeepDives.AutogradTransforms.Correctness.fourthPower_vjp (s : Spec.Shape) (inputs gradient : TorchLean.TensorPack ℝ [s]) (seed value : TorchLean.Tensor ℝ s) (checked : Runtime.Autograd.Torch.TypedGraphWithData.vjpChecked ((fourthPower s).toTypedGraph (output s)) inputs () seed = Except.ok (gradient, value)) : Proofs.Autograd.flattenCtx gradient = (ContinuousLinearMap.adjoint (fderiv ℝ (fun (x : Proofs.Autograd.CtxVec [s]) => Proofs.Autograd.tensorToVec (((fourthPower s).toTypedGraph (output s)).forward (Proofs.Autograd.unflattenCtx x))) (Proofs.Autograd.flattenCtx inputs))) (Proofs.Autograd.tensorToVec seed)

A successful checked reverse pass computes the adjoint of the forward function's derivative.

This theorem is about the exact-real graph. Running the same arithmetic in floating point needs a separate rounding-error argument, and higher derivatives need more than this first-order proof.

def NN.Examples.DeepDives.AutogradTransforms.HigherOrder.expSquare {α : Type} [TorchLean.Storage α] [Context α] {shape : Spec.Shape} (x : TorchLean.Tensor α shape) : TorchLean.Tensor α shape

Apply exp(x*x) to a tensor, using the caller's scalar backend and shape.

theorem NN.Examples.DeepDives.AutogradTransforms.HigherOrder.expSquare_derivative {shape : Spec.Shape} (n : ℕ) (directions : Fin n → TorchLean.Tensor ℝ shape) (x : TorchLean.Tensor ℝ shape) : Runtime.Autograd.Model.Dual.Nested.tangentTensor (expSquare (Runtime.Autograd.Model.Dual.Nested.seedTensor directions x)) = (iteratedFDeriv ℝ n (fun (y : TorchLean.Tensor ℝ shape) => TorchLean.Tensor.map (fun (z : ℝ) => Real.exp (z * z)) y) x) directions

Every derivative order, tensor shape, and direction tuple share one proof.

The seed and extraction functions are also used by autograd.model.derivative. This theorem checks the tensor computation between them; model lowering requires its own correctness evidence.

def NN.Examples.DeepDives.AutogradTransforms.HigherOrder.expSquareProgram {α : Type} [TorchLean.Storage α] [Context α] (shape : Spec.Shape) : Runtime.Autograd.TypedGraph.GraphM.M α [shape] (Runtime.Autograd.TypedGraph.GraphM.Var shape)

Record squaring and exponentiation as two graph operations, for any scalar backend.

theorem NN.Examples.DeepDives.AutogradTransforms.HigherOrder.expSquareProgram_eval {α : Type} [TorchLean.Storage α] [Context α] {shape : Spec.Shape} (x : TorchLean.Tensor α shape) : Except.map (fun (graph : Runtime.Autograd.Torch.TypedGraph α [shape] shape) => graph.forward (TorchLean.TensorPack.singleton x)) (Runtime.Autograd.Torch.lowerToTypedGraph (expSquareProgram shape)) = Except.ok (x.mulSpec x).expSpec

Recording succeeds and the resulting graph evaluates the two requested operations.

The proof keeps the scalar backend abstract: recording checks shapes and indices, not numerical values. This also prevents the proof from expanding real or nested-dual arithmetic internals.

theorem NN.Examples.DeepDives.AutogradTransforms.HigherOrder.expSquareProgram_derivative {shape : Spec.Shape} (n : ℕ) (real : Runtime.Autograd.Torch.TypedGraph ℝ [shape] shape) (nested : Runtime.Autograd.Torch.TypedGraph (Runtime.Autograd.Model.Dual.Nested ℝ n) [shape] shape) (hr : Runtime.Autograd.Torch.lowerToTypedGraph (expSquareProgram shape) = Except.ok real) (hn : Runtime.Autograd.Torch.lowerToTypedGraph (expSquareProgram shape) = Except.ok nested) (directions : Fin n → TorchLean.Tensor ℝ shape) (x : TorchLean.Tensor ℝ shape) : Runtime.Autograd.Model.Dual.Nested.tangentTensor (nested.forward (TorchLean.TensorPack.singleton (Runtime.Autograd.Model.Dual.Nested.seedTensor directions x))) = (iteratedFDeriv ℝ n (fun (y : TorchLean.Tensor ℝ shape) => real.forward (TorchLean.TensorPack.singleton y)) x) directions

Nested-dual execution of the recorded graph computes every mixed derivative of its real run.

Unlike the tensor-only example above, this statement includes the actual recorder and graph execution. It does not assert correctness of arbitrary model lowering or the reverse pass.

def NN.Examples.DeepDives.AutogradTransforms.HigherOrder.squareModel (shape : Spec.Shape) : TorchLean.nn.Sequential shape shape

A parameter-free model using the same square layer as executable neural networks.

theorem NN.Examples.DeepDives.AutogradTransforms.HigherOrder.squareModel_derivative {shape : Spec.Shape} (x : TorchLean.Tensor ℝ shape) (directions : List (TorchLean.Tensor ℝ shape)) : TorchLean.autograd.model.derivative (squareModel shape) TorchLean.nn.State.empty x directions = pure ((iteratedFDeriv ℝ directions.length (fun (y : TorchLean.Tensor ℝ shape) => y.mulSpec y) x) fun (i : Fin directions.length) => directions[i])

The public IO transform returns the exact mixed derivative, for every shape and order.

There is no assumed lowering result: the proof checks model validation, graph recording, and execution before applying autograd to the arithmetic. The empty direction list is included.