Automatic Differentiation

This module contains gradient, VJP, Jacobian, JVP, and HVP operations for models and pure one-argument tensor functions.

Autograd operations (grad/vjp/jacobian) over TorchLean programs.

This namespace is conceptually similar to PyTorch autograd + functorch/torch.func:

• gradients of losses w.r.t. parameters and inputs
• VJPs and Jacobians for analysis and verification tooling

PyTorch references:

• Autograd: https://pytorch.org/docs/stable/autograd.html
• torch.func (jacfwd/jacrev, etc.): https://pytorch.org/docs/stable/func.html

@[reducible, inline]

abbrev TorchLean.autograd.model.Params {σ τ : Spec.Shape} (model : nn.Sequential σ τ) (α : Type) : Type

Parameter tensors for model, indexed by its statically known parameter shapes.

@[reducible, inline]

abbrev TorchLean.autograd.model.OutputLoss (τ υ : Spec.Shape) : Type 1

A scalar loss computed from a model output and its target.

def TorchLean.autograd.model.initParamsWith {σ τ : Spec.Shape} (model : nn.Sequential σ τ) {α : Type} (cast : Float → α) : Params model α

Cast a model's initial Float parameters into another scalar representation.

def TorchLean.autograd.model.initParams {σ τ : Spec.Shape} (model : nn.Sequential σ τ) {α : Type} [Runtime.FromFloat α] : Params model α

Initialize model parameters in a scalar representation that accepts host Float values.

def TorchLean.autograd.model.OutputLoss.mse {τ : Spec.Shape} (reduction : Loss.Reduction := Loss.Reduction.mean) : OutputLoss τ τ

Mean-squared error between a model output and its target.

def TorchLean.autograd.model.OutputLoss.crossEntropyOneHot {τ : Spec.Shape} (reduction : Loss.Reduction := Loss.Reduction.mean) : OutputLoss τ τ

Cross-entropy between logits and one-hot targets.

def TorchLean.autograd.model.OutputLoss.detach {τ υ : Spec.Shape} (loss : OutputLoss τ υ) : OutputLoss τ υ

Stop gradients through the model output before evaluating loss.

def TorchLean.autograd.model.lossProgram {σ τ υ : Spec.Shape} (model : nn.Sequential σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] : Runtime.Autograd.TorchLean.Program α (Runtime.Autograd.TorchLean.NN.Seq.paramShapes model ++ [σ, υ]) Spec.Shape.scalar

Compile loss (model params input) target as a scalar TorchLean program.

def TorchLean.autograd.model.gradParams {σ τ υ : Spec.Shape} (model : nn.Sequential σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (Params model α)

Gradient of a model-loss w.r.t. the model parameters.

Common training use case. PyTorch analogue: loss.backward() followed by parameter updates.

def TorchLean.autograd.model.gradInputs {σ τ υ : Spec.Shape} (model : nn.Sequential σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (TensorPack α [σ, υ])

Gradient of the loss w.r.t. the inputs (x and target).

def TorchLean.autograd.model.gradX {σ τ υ : Spec.Shape} (model : nn.Sequential σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (Spec.Tensor α σ)

Convenience: gradient of the loss w.r.t. x.

def TorchLean.autograd.model.gradTarget {σ τ υ : Spec.Shape} (model : nn.Sequential σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (Spec.Tensor α υ)

Convenience: gradient of the loss w.r.t. the target argument.

structure TorchLean.autograd.model.ValueAndGrads {σ τ υ : Spec.Shape} (model : nn.Sequential σ τ) (α : Type) : Type

Forward+backward result for a scalar loss built from a model output.

PyTorch comparison: this is the "compute loss + backward" payload, but with shapes tracked.

• value : Spec.Tensor α Spec.Shape.scalar

  Value at the current point.

• dparams : Params model α

  Gradients w.r.t. parameters.

• dx : Spec.Tensor α σ

  Gradient w.r.t. input.

• dtarget : Spec.Tensor α υ

  Gradient w.r.t. target.

def TorchLean.autograd.model.valueAndGrads {σ τ υ : Spec.Shape} (model : nn.Sequential σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (ValueAndGrads model α)

Run loss(model(params, x), target) and compute gradients w.r.t:

• model parameters,
• x,
• target.

This hides the CompiledScalar/argument-pack boilerplate for the common "one sample" case.

def TorchLean.autograd.model.valueAndGradParams {σ τ υ : Spec.Shape} (model : nn.Sequential σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (Spec.Tensor α Spec.Shape.scalar × Params model α)

Return the scalar loss tensor together with gradients for the model parameters.

def TorchLean.autograd.model.valueAndGradParamsScalar {σ τ υ : Spec.Shape} (model : nn.Sequential σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (α × Params model α)

valueAndGradParams, but convert the 0-dim loss tensor to a scalar α.

def TorchLean.autograd.model.valueAndGradX {σ τ υ : Spec.Shape} (model : nn.Sequential σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (Spec.Tensor α Spec.Shape.scalar × Spec.Tensor α σ)

Return (loss_value, grad_x).

def TorchLean.autograd.model.valueAndGradTarget {σ τ υ : Spec.Shape} (model : nn.Sequential σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (Spec.Tensor α Spec.Shape.scalar × Spec.Tensor α υ)

Return (loss_value, grad_target).

def TorchLean.autograd.model.vjpParams {σ τ : Spec.Shape} (model : nn.Sequential σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (seedOut : Spec.Tensor α τ) : IO (Params model α)

Vector-Jacobian product (VJP) w.r.t. model parameters.

Primitive for sending output cotangents back into parameters. Use it for custom losses or analysis tooling when you already have a seed tensor seedOut : τ.

def TorchLean.autograd.model.vjpInputs {σ τ : Spec.Shape} (model : nn.Sequential σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (seedOut : Spec.Tensor α τ) : IO (TensorPack α [σ])

VJP w.r.t. the model input.

Returns a one-element _root_.TorchLean.TensorPack to match the general "inputs list" API shape. For the common case, use vjpInput to get the tensor directly.

def TorchLean.autograd.model.vjpInput {σ τ : Spec.Shape} (model : nn.Sequential σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (seedOut : Spec.Tensor α τ) : IO (Spec.Tensor α σ)

Vector-Jacobian product with respect to the single model input tensor.

def TorchLean.autograd.model.jacrevParams {σ τ : Spec.Shape} (model : nn.Sequential σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) : IO (Array (Params model α))

Reverse-mode Jacobian (jacrev) of the model output w.r.t. parameters.

Returns an array of parameter-structured gradients: one entry per output coordinate. This mirrors the usual "jacrev returns a stack of per-output gradients" shape.

def TorchLean.autograd.model.jvpParams {σ τ υ : Spec.Shape} (model : nn.Sequential σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) (vparams : Params model α) : IO α

Jacobian-vector product (JVP) of a scalar loss w.r.t. parameters.

Directional derivative in the direction vparams. Conceptually:

$$ \left.\frac{d}{dt} \operatorname{loss}(\mathrm{params}+t\,\mathrm{vparams},x,\mathrm{target}) \right|_{t=0}. $$

def TorchLean.autograd.model.hvpParams {σ τ υ : Spec.Shape} (model : nn.Sequential σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) (vparams : Params model α) : IO (Params model α)

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

Returns a parameter-structured tensor list of the same shape as params.

In PyTorch terms, this is the "functorch" style: differentiate plain functions, not modules.

@[reducible, inline]

abbrev TorchLean.autograd.func.Fn (σ τ : Spec.Shape) : Type 1

A scalar-polymorphic tensor function written against TorchLean's differentiable operations.

def TorchLean.autograd.func.forwardProgram {σ τ : Spec.Shape} (f : Fn σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] : Runtime.Autograd.TorchLean.Program α [σ] τ

Adapt a tensor function to the single-input program representation used by autograd.

def TorchLean.autograd.func.jacfwd {σ τ : Spec.Shape} (f : Fn σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] (x : Spec.Tensor α σ) : IO (Array (Spec.Tensor α τ))

Forward-mode Jacobian (jacfwd) for a pure tensor function.

def TorchLean.autograd.func.hessian {σ : Spec.Shape} (f : Fn σ Spec.Shape.scalar) {α : Type} [Context α] [DecidableEq Spec.Shape] (x : Spec.Tensor α σ) : IO (Array (Spec.Tensor α σ))

Hessian for a scalar-valued function.

def TorchLean.autograd.func.vjp {σ τ : Spec.Shape} (f : Fn σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] (x : Spec.Tensor α σ) (seedOut : Spec.Tensor α τ) : IO (Spec.Tensor α σ)

Vector-Jacobian product (VJP) for a pure function.

def TorchLean.autograd.func.jacrev {σ τ : Spec.Shape} (f : Fn σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] (x : Spec.Tensor α σ) : IO (Array (Spec.Tensor α σ))

Reverse-mode Jacobian (jacrev) of a pure tensor function.

Returns the Jacobian rows as an array of doutput/dinput tensors.

def TorchLean.autograd.func.grad {σ : Spec.Shape} (f : Fn σ Spec.Shape.scalar) {α : Type} [Context α] [DecidableEq Spec.Shape] (x : Spec.Tensor α σ) : IO (Spec.Tensor α σ)

Gradient of a scalar-valued function w.r.t. its input.

def TorchLean.autograd.func.valueAndGrad {σ : Spec.Shape} (f : Fn σ Spec.Shape.scalar) {α : Type} [Context α] [DecidableEq Spec.Shape] (x : Spec.Tensor α σ) : IO (Spec.Tensor α Spec.Shape.scalar × Spec.Tensor α σ)

Return (value, grad) for a scalar-valued function at x.

def TorchLean.autograd.func.valueAndGradScalar {σ : Spec.Shape} (f : Fn σ Spec.Shape.scalar) {α : Type} [Context α] [DecidableEq Spec.Shape] (x : Spec.Tensor α σ) : IO (α × Spec.Tensor α σ)

valueAndGrad, but convert the 0-dim value tensor to a scalar α.