Model Automatic Differentiation

Model state, output losses, and differentiation with respect to model state and inputs. Import NN.API.Autograd for the complete public autograd API.

@[reducible, inline]

abbrev TorchLean.autograd.model.State {σ τ : Shape} (model : nn.Sequential σ τ) (α : Type) [Storage α] : Type

Complete model state, indexed by its statically known tensor shapes.

def TorchLean.autograd.model.fullState {σ τ : Shape} (model : nn.Sequential σ τ) {α : Type} [Storage α] (value : α) : State model α

Construct a model-shaped state whose every tensor contains value.

structure TorchLean.autograd.model.Loss (τ υ : Shape) : Type 1

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

• forward {α : Type} [Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] : Runtime.ValueRef m α τ → Runtime.ValueRef m α υ → m (Runtime.ValueRef m α [])

  Operation-polymorphic loss program.

• validate : Except String Unit

  Configuration checks performed before lowering or execution.

def TorchLean.autograd.model.initialState {σ τ : Shape} (model : nn.Sequential σ τ) {α : Type} [Storage α] [Runtime.FromFloat α] : State model α

Initialize model state in an element type that accepts host Float values.

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

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

def TorchLean.autograd.model.Loss.oneHotCrossEntropy {τ : Shape} (axis : ℕ) (reduction : Loss.Reduction := Loss.Reduction.mean) : Loss τ τ

Cross-entropy between logits and one-hot targets along the selected class dimension.

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

Stop gradients through the model output before evaluating loss.

def TorchLean.autograd.model.Internal.lossProgram {σ τ υ : Shape} (model : nn.Sequential σ τ) (loss : Loss τ υ) {α : Type} [Storage α] [Context α] : Runtime.Autograd.Model.Program α (Runtime.Autograd.Model.Layers.Seq.stateShapes model ++ [σ, υ]) []

Lower loss (model state input) target to the typed scalar program used by autograd.

def TorchLean.autograd.model.Internal.validateLoss {σ τ υ : Shape} (model : nn.Sequential σ τ) (loss : Loss τ υ) : IO Unit

Reject an invalid model or loss before lowering an autograd program.

def TorchLean.autograd.model.grad {σ τ υ : Shape} (model : nn.Sequential σ τ) (loss : Loss τ υ) {α : Type} [Storage α] [Context α] (state : State model α) (input : Tensor α σ) (target : Tensor α υ) (value : Bool := false) : IO (match value with | false => State model α | true => State model α × Tensor α [])

Differentiate a model loss with respect to every tensor in the model state.

The result has the same shape-indexed layout as State model α. Set value := true to return (grad, lossValue) from the same evaluation. For models with persistent buffers, this computes mathematical sensitivities for those entries as well; nn.requiresGrad separately controls which state tensors an optimizer updates.

def TorchLean.autograd.model.vjp {σ τ : Shape} (model : nn.Sequential σ τ) {α : Type} [Storage α] [Context α] (state : State model α) (input : Tensor α σ) (outputGradient : Tensor α τ) : IO (State model α × Tensor α σ)

Vector-Jacobian product with respect to the model and its input.

The returned pair is (stateGrad, inputGrad). Both values come from one reverse pass.

def TorchLean.autograd.model.jacrev {σ τ : Shape} (model : nn.Sequential σ τ) {α : Type} [Storage α] [Context α] (state : State model α) (input : Tensor α σ) : IO (nn.State α (List.map τ.concat (Runtime.Autograd.Model.Layers.Seq.stateShapes model)))

Reverse-mode Jacobian (jacrev) of the model output with respect to model state.

Returns one Jacobian tensor per state tensor. Each has the output axes followed by that state tensor's axes, so parameter shapes stay distinct without an outer array of gradient states.

def TorchLean.autograd.model.jvp {σ τ υ : Shape} (model : nn.Sequential σ τ) (loss : Loss τ υ) {α : Type} [Storage α] [Context α] (state : State model α) (input : Tensor α σ) (target : Tensor α υ) (stateDirection : State model α) : IO (Tensor α [])

Jacobian-vector product (JVP) of a scalar loss with respect to model state.

Directional derivative in the direction stateDirection. Conceptually:

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

def TorchLean.autograd.model.hvp {σ τ υ : Shape} (model : nn.Sequential σ τ) (loss : Loss τ υ) {α : Type} [Storage α] [Context α] (state : State model α) (input : Tensor α σ) (target : Tensor α υ) (stateDirection : State model α) : IO (State model α)

Hessian-vector product (HVP) of a scalar loss with respect to model state.

Returns model state with the same shape layout as state.