Runtime Autograd Helpers

Model-shaped and function-shaped differentiation helpers over TorchLean runtime programs.

@[reducible, inline]

abbrev NN.API.TorchLean.Autodiff.Model.Params {σ τ : Spec.Shape} (model : NN.Seq σ τ) (α : Type) : Type

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

@[reducible, inline]

abbrev NN.API.TorchLean.Autodiff.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.

def NN.API.TorchLean.Autodiff.Model.initParamsWith {σ τ : Spec.Shape} (model : NN.Seq σ τ) {α : Type} (cast : Float → α) : Params model α

Initialize model parameters by casting the model's Float initializers elementwise using cast.

def NN.API.TorchLean.Autodiff.Model.initParams {σ τ : Spec.Shape} (model : NN.Seq σ τ) {α : Type} [Runtime.Scalar α] : Params model α

Initialize model parameters using the runtime literal injection API.Runtime.ofFloat.

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

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

def NN.API.TorchLean.Autodiff.Model.OutputLoss.mse {τ : Spec.Shape} (reduction : Loss.Reduction := Runtime.Autograd.TorchLean.Loss.Reduction.mean) : OutputLoss τ τ

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

def NN.API.TorchLean.Autodiff.Model.OutputLoss.crossEntropyOneHot {τ : Spec.Shape} (reduction : Loss.Reduction := Runtime.Autograd.TorchLean.Loss.Reduction.mean) : OutputLoss τ τ

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

def NN.API.TorchLean.Autodiff.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.TorchLean.Autodiff.Model.lossProgram {σ τ υ : Spec.Shape} (model : NN.Seq σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] : Program α (model.paramShapes ++ [σ, υ]) Spec.Shape.scalar

Build a TorchLean Program that computes a scalar loss from (params, x, target).

This is the bridge between Seq.program (which produces model outputs) and the autograd entry points (which expect a scalar-valued program).

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

VJP of the model output w.r.t. parameters.

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

VJP of the model output w.r.t. inputs.

def NN.API.TorchLean.Autodiff.Model.jacrevParams {σ τ : Spec.Shape} (model : NN.Seq σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) : IO (Array (Params model α))

Jacobian (reverse-mode) of the model output w.r.t. parameters, returned as rows.

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

Gradient of loss(model(params, x), target) w.r.t. parameters.

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

Gradient of loss(model(params, x), target) w.r.t. inputs (x and target).

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

JVP of a scalar loss w.r.t. parameters in direction vparams.

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

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

@[reducible, inline]

abbrev NN.API.TorchLean.Autodiff.Function1.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.TorchLean.Autodiff.Function1.program {σ τ : Spec.Shape} (f : Fn σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] : Program α [σ] τ

Turn an Fn into a single-input TorchLean Program.

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

Forward-mode Jacobian (rows) of a pure function.

def NN.API.TorchLean.Autodiff.Function1.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 pure function.