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.Internal.pullback {α : 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)

Return the seeded input gradients and forward result, reporting domain errors through IO.

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. Scalar outputs use the same path with shape [].

def Runtime.Autograd.Model.Autodiff.jacrevParams {α : 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.jacrevInputs {α : 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.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.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.

This is the general VJP with a unit scalar seed.

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.Internal.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.