Autodiff

Autodiff utilities beyond basic .backward():

• hvpParams: Hessian-vector product for scalar losses w.r.t. parameters, using forward-over-reverse via Dual scalars.

This is runtime/executable functionality intended for TorchLean ergonomics; it is separate from the fderiv proof developments.

def Runtime.Autograd.Model.Autodiff.Impl.okOrThrow {α : Type} : Result α → IO α

Unwrap a runtime Result into IO, throwing a user error on failure.

This is used throughout this module because lowering and backpropagation utilities return an Autograd.Result with a structured error message.

def Runtime.Autograd.Model.Autodiff.Impl.vjpWithValue {α : Type} [TorchLean.Storage α] [Add α] [Zero α] {shapes : List Spec.Shape} {output : Spec.Shape} (graph : Torch.TypedGraph α shapes output) (inputs : TorchLean.TensorPack α shapes) (seed : TorchLean.Tensor α output) : IO (TorchLean.TensorPack α shapes × TorchLean.Tensor α output)

Execute the checked graph pullback, reporting domain errors through the IO API.

def Runtime.Autograd.Model.Autodiff.lowerScalarToTypedGraph {α : Type} [TorchLean.Storage α] [Context α] {paramShapes inputShapes : List Spec.Shape} (program : {β : Type} → [inst : TorchLean.Storage β] → [inst_1 : Context β] → Program β (paramShapes ++ inputShapes) Spec.Shape.scalar) : IO (Torch.TypedScalarGraph α (paramShapes ++ inputShapes))

Lower a scalar-valued TorchLean program to a reusable typed graph.

def Runtime.Autograd.Model.Autodiff.lowerToTypedGraph {α : Type} [TorchLean.Storage α] [Context α] {paramShapes inputShapes : List Spec.Shape} {τ : Spec.Shape} (f : {β : Type} → [inst : TorchLean.Storage β] → [inst_1 : Context β] → Program β (paramShapes ++ inputShapes) τ) : IO (Torch.TypedGraph α (paramShapes ++ inputShapes) τ)

Lower a TorchLean program to a reusable TypedGraph.

The graph retains a typed reference to the value returned by the program, which may be an input or any recorded node. This is the tensor-output analogue of lowerScalarToTypedGraph; it is used by jacrevOut* and vjpOut*.

def Runtime.Autograd.Model.Autodiff.jacrevOutParams {α : Type} [TorchLean.Storage α] [Context α] {paramShapes inputShapes : List Spec.Shape} {τ : Spec.Shape} (f : {β : Type} → [inst : TorchLean.Storage β] → [inst_1 : Context β] → Program β (paramShapes ++ inputShapes) τ) (params : TorchLean.TensorPack α paramShapes) (xs : TorchLean.TensorPack α inputShapes) : IO (TorchLean.TensorPack α (List.map τ.concat paramShapes))

Reverse Jacobian, with output axes prepended to each parameter tensor.

def Runtime.Autograd.Model.Autodiff.jacrevOutInputs {α : Type} [TorchLean.Storage α] [Context α] {paramShapes inputShapes : List Spec.Shape} {τ : Spec.Shape} (f : {β : Type} → [inst : TorchLean.Storage β] → [inst_1 : Context β] → Program β (paramShapes ++ inputShapes) τ) (params : TorchLean.TensorPack α paramShapes) (xs : TorchLean.TensorPack α inputShapes) : IO (TorchLean.TensorPack α (List.map τ.concat inputShapes))

Reverse Jacobian, with output axes prepended to each input tensor.

def Runtime.Autograd.Model.Autodiff.jacfwdInput {α : Type} [TorchLean.Storage α] [Context α] {σ τ : Spec.Shape} (f : {β : Type} → [inst : TorchLean.Storage β] → [inst_1 : Context β] → Program β [σ] τ) (x : TorchLean.Tensor α σ) : IO (TorchLean.Tensor α (τ.concat σ))

Compute the forward Jacobian for a single tensor input.

Each input basis vector runs through the checked graph JVP. This uses the same primitive rules as directional differentiation, including stopped gradients and the runtime's choices at nonsmooth points. The returned tensor has output axes followed by input axes.

def Runtime.Autograd.Model.Autodiff.gradients {α : Type} [TorchLean.Storage α] [Context α] {paramShapes inputShapes : List Spec.Shape} (loss : {β : Type} → [inst : TorchLean.Storage β] → [inst_1 : Context β] → Program β (paramShapes ++ inputShapes) Spec.Shape.scalar) (params : TorchLean.TensorPack α paramShapes) (xs : TorchLean.TensorPack α inputShapes) : IO (TorchLean.TensorPack α paramShapes × TorchLean.TensorPack α inputShapes)

Differentiate a scalar loss with respect to parameters and inputs in one reverse pass.

def Runtime.Autograd.Model.Autodiff.vjp {α : Type} [TorchLean.Storage α] [Context α] {paramShapes inputShapes : List Spec.Shape} {τ : Spec.Shape} (f : {β : Type} → [inst : TorchLean.Storage β] → [inst_1 : Context β] → Program β (paramShapes ++ inputShapes) τ) (params : TorchLean.TensorPack α paramShapes) (xs : TorchLean.TensorPack α inputShapes) (seedOut : TorchLean.Tensor α τ) : IO (TorchLean.TensorPack α paramShapes × TorchLean.TensorPack α inputShapes)

Compute a tensor-output VJP with respect to parameters and inputs in one reverse pass.

def Runtime.Autograd.Model.Autodiff.jvpLossParams {α : Type} [TorchLean.Storage α] [Context α] {paramShapes inputShapes : List Spec.Shape} (loss : {β : Type} → [inst : TorchLean.Storage β] → [inst_1 : Context β] → Program β (paramShapes ++ inputShapes) Spec.Shape.scalar) (params : TorchLean.TensorPack α paramShapes) (xs : TorchLean.TensorPack α inputShapes) (vparams : TorchLean.TensorPack α paramShapes) : IO (TorchLean.Tensor α [])

Directional derivative of scalar loss along vparams (forward-mode JVP).

def Runtime.Autograd.Model.Autodiff.Impl.dualGradients {α : Type} [TorchLean.Storage α] [Context α] {paramShapes inputShapes : List Spec.Shape} (loss : {β : Type} → [inst : TorchLean.Storage β] → [inst_1 : Context β] → Program β (paramShapes ++ inputShapes) Spec.Shape.scalar) (argsD : TorchLean.TensorPack (Dual α) (paramShapes ++ inputShapes)) : IO (TorchLean.TensorPack (Dual α) (paramShapes ++ inputShapes))

Reverse-mode gradients over dual-valued arguments, shared by both HVP projections.

def Runtime.Autograd.Model.Autodiff.hvpParams {α : Type} [TorchLean.Storage α] [Context α] {paramShapes inputShapes : List Spec.Shape} (loss : {β : Type} → [inst : TorchLean.Storage β] → [inst_1 : Context β] → Program β (paramShapes ++ inputShapes) Spec.Shape.scalar) (params : TorchLean.TensorPack α paramShapes) (xs : TorchLean.TensorPack α inputShapes) (vparams : TorchLean.TensorPack α paramShapes) : IO (TorchLean.TensorPack α paramShapes)

Hessian-vector product (HVP) for a scalar loss w.r.t. parameters.

This computes d/dε (∇_params loss(params + ε*vparams)) |_{ε=0} and returns a TorchLean.TensorPack aligned with paramShapes.

Implementation: run reverse-mode AD over dual scalars (Dual), with parameter tangents set to vparams and input tangents set to 0. The tangent part of the resulting gradients is the HVP.

def Runtime.Autograd.Model.Autodiff.hvpInputs {α : Type} [TorchLean.Storage α] [Context α] {paramShapes inputShapes : List Spec.Shape} (loss : {β : Type} → [inst : TorchLean.Storage β] → [inst_1 : Context β] → Program β (paramShapes ++ inputShapes) Spec.Shape.scalar) (params : TorchLean.TensorPack α paramShapes) (xs vxs : TorchLean.TensorPack α inputShapes) : IO (TorchLean.TensorPack α inputShapes)

Hessian-vector product (HVP) for a scalar loss w.r.t. inputs.

This computes d/dε (∇_xs loss(xs + ε*vxs)) |_{ε=0} and returns a TorchLean.TensorPack aligned with inputShapes.

Implementation: the same forward-over-reverse trick as hvpParams, but we attach tangents to inputs instead of parameters.

def Runtime.Autograd.Model.Autodiff.hessianInput {α : Type} [TorchLean.Storage α] [Context α] {σ : Spec.Shape} (f : {β : Type} → [inst : TorchLean.Storage β] → [inst_1 : Context β] → Program β [σ] Spec.Shape.scalar) (x : TorchLean.Tensor α σ) : IO (TorchLean.Tensor α (σ.concat σ))

Full Hessian tensor for a scalar function of a single tensor input.

Columns are evaluated as H * e_i in the flattened input basis, then arranged with one copy of the input axes for each derivative.