Public autograd operations

This module contains the public gradient, VJP, Jacobian, JVP, and HVP APIs 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 NN.API.autograd.model.Params {σ τ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) (α : Type) : Type

Parameter pack type for a given model (a TensorPack over Seq.paramShapes).

@[reducible, inline]

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

Loss function over a model output and a target.

This is expressed in terms of RefTy so it works uniformly for eager execution and compiled execution.

@[reducible, inline]

abbrev NN.API.autograd.model.linearParams {α : Type} {inDim outDim seedW seedB : ℕ} (w : Spec.Tensor α (Tensor.Shape.Mat outDim inDim)) (b : Spec.Tensor α (Tensor.Shape.Vec outDim)) : Params (TorchLean.Layers.linear inDim outDim seedW seedB) α

Pack explicit weight and bias tensors for a single Layers.linear model.

@[reducible, inline]

abbrev NN.API.autograd.model.OutputLoss.mse {τ : Spec.Shape} (reduction : TorchLean.Loss.Reduction := Runtime.Autograd.TorchLean.Loss.Reduction.mean) : OutputLoss τ τ

Mean-squared error loss (mse) between yhat and y.

@[reducible, inline]

abbrev NN.API.autograd.model.OutputLoss.crossEntropyOneHot {τ : Spec.Shape} (reduction : TorchLean.Loss.Reduction := Runtime.Autograd.TorchLean.Loss.Reduction.mean) : OutputLoss τ τ

Cross-entropy loss between logits and one-hot targets. PyTorch analogue: nn.CrossEntropyLoss.

@[reducible, inline]

abbrev NN.API.autograd.model.OutputLoss.detach {τ υ : Spec.Shape} (loss : OutputLoss τ υ) : OutputLoss τ υ

Detach the model output before feeding it into a loss.

This is useful when you want to compute a metric loss without backpropagating through it.

def NN.API.autograd.model.gradParams {σ τ υ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) (loss : TorchLean.Autodiff.Model.OutputLoss τ υ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (params : TorchLean.Autodiff.Model.Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (TorchLean.Autodiff.Model.Params model α)

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

This is the common training use case (PyTorch analogue: loss.backward() followed by parameter updates).

def NN.API.autograd.model.gradInputs {σ τ υ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) (loss : TorchLean.Autodiff.Model.OutputLoss τ υ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (params : TorchLean.Autodiff.Model.Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (TensorPack α [σ, υ])

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

def NN.API.autograd.model.gradX {σ τ υ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) (loss : TorchLean.Autodiff.Model.OutputLoss τ υ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (params : TorchLean.Autodiff.Model.Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (Spec.Tensor α σ)

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

def NN.API.autograd.model.gradTarget {σ τ υ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) (loss : TorchLean.Autodiff.Model.OutputLoss τ υ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (params : TorchLean.Autodiff.Model.Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (Spec.Tensor α υ)

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

structure NN.API.autograd.model.ValueAndGrads {σ τ υ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) (α : 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 : TorchLean.Autodiff.Model.Params model α

  Gradients w.r.t. parameters.

• dx : Spec.Tensor α σ

  Gradient w.r.t. input.

• dtarget : Spec.Tensor α υ

  Gradient w.r.t. target.

def NN.API.autograd.model.valueAndGrads {σ τ υ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) (loss : TorchLean.Autodiff.Model.OutputLoss τ υ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (params : TorchLean.Autodiff.Model.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 NN.API.autograd.model.valueAndGradParams {σ τ υ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) (loss : TorchLean.Autodiff.Model.OutputLoss τ υ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (params : TorchLean.Autodiff.Model.Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (Spec.Tensor α Spec.Shape.scalar × TorchLean.Autodiff.Model.Params model α)

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

def NN.API.autograd.model.valueAndGradParamsScalar {σ τ υ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) (loss : TorchLean.Autodiff.Model.OutputLoss τ υ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (params : TorchLean.Autodiff.Model.Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (α × TorchLean.Autodiff.Model.Params model α)

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

def NN.API.autograd.model.valueAndGradX {σ τ υ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) (loss : TorchLean.Autodiff.Model.OutputLoss τ υ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (params : TorchLean.Autodiff.Model.Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (Spec.Tensor α Spec.Shape.scalar × Spec.Tensor α σ)

Return (loss_value, grad_x).

def NN.API.autograd.model.valueAndGradTarget {σ τ υ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) (loss : TorchLean.Autodiff.Model.OutputLoss τ υ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (params : TorchLean.Autodiff.Model.Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (Spec.Tensor α Spec.Shape.scalar × Spec.Tensor α υ)

Return (loss_value, grad_target).

def NN.API.autograd.model.vjpParams {σ τ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (params : TorchLean.Autodiff.Model.Params model α) (x : Spec.Tensor α σ) (seedOut : Spec.Tensor α τ) : IO (TorchLean.Autodiff.Model.Params model α)

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

This is the "grad of outputs back into parameters" primitive. It is useful for custom losses or analysis tooling when you already have a seed tensor seedOut : τ.

def NN.API.autograd.model.vjpInputs {σ τ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (params : TorchLean.Autodiff.Model.Params model α) (x : Spec.Tensor α σ) (seedOut : Spec.Tensor α τ) : IO (TensorPack α [σ])

VJP w.r.t. the model input.

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

def NN.API.autograd.model.vjpInput {σ τ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (params : TorchLean.Autodiff.Model.Params model α) (x : Spec.Tensor α σ) (seedOut : Spec.Tensor α τ) : IO (Spec.Tensor α σ)

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

def NN.API.autograd.model.jacrevParams {σ τ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (params : TorchLean.Autodiff.Model.Params model α) (x : Spec.Tensor α σ) : IO (Array (TorchLean.Autodiff.Model.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 NN.API.autograd.model.jvpParams {σ τ υ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) (loss : TorchLean.Autodiff.Model.OutputLoss τ υ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (params : TorchLean.Autodiff.Model.Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) (vparams : TorchLean.Autodiff.Model.Params model α) : IO α

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

This is the directional derivative in the direction vparams. Conceptually: d/dt loss(params + t*vparams, x, target) | t = 0.

def NN.API.autograd.model.hvpParams {σ τ υ : Spec.Shape} (model : TorchLean.NN.Seq σ τ) (loss : TorchLean.Autodiff.Model.OutputLoss τ υ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (params : TorchLean.Autodiff.Model.Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) (vparams : TorchLean.Autodiff.Model.Params model α) : IO (TorchLean.Autodiff.Model.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 NN.API.autograd.fn1.Fn (σ τ : Spec.Shape) : Type 1

Type of a pure tensor function expressed in RefTy form.

This matches the calling convention expected by TorchLean.Program/autodiff compilation.

def NN.API.autograd.fn1.jacfwd {σ τ : Spec.Shape} (f : TorchLean.Autodiff.Function1.Fn σ τ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (x : Spec.Tensor α σ) : IO (Array (Spec.Tensor α τ))

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

def NN.API.autograd.fn1.hessian {σ : Spec.Shape} (f : TorchLean.Autodiff.Function1.Fn σ Spec.Shape.scalar) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (x : Spec.Tensor α σ) : IO (Array (Spec.Tensor α σ))

Hessian for a scalar-valued function.

def NN.API.autograd.fn1.vjp {σ τ : Spec.Shape} (f : TorchLean.Autodiff.Function1.Fn σ τ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (x : Spec.Tensor α σ) (seedOut : Spec.Tensor α τ) : IO (Spec.Tensor α σ)

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

def NN.API.autograd.fn1.jacrev {σ τ : Spec.Shape} (f : TorchLean.Autodiff.Function1.Fn σ τ) {α : Type} [Semantics.Scalar α] [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 NN.API.autograd.fn1.grad {σ : Spec.Shape} (f : TorchLean.Autodiff.Function1.Fn σ Spec.Shape.scalar) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (x : Spec.Tensor α σ) : IO (Spec.Tensor α σ)

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

def NN.API.autograd.fn1.valueAndGrad {σ : Spec.Shape} (f : TorchLean.Autodiff.Function1.Fn σ Spec.Shape.scalar) {α : Type} [Semantics.Scalar α] [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 NN.API.autograd.fn1.valueAndGradScalar {σ : Spec.Shape} (f : TorchLean.Autodiff.Function1.Fn σ Spec.Shape.scalar) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (x : Spec.Tensor α σ) : IO (α × Spec.Tensor α σ)

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