Backward Passes and Optimizers

Gradient extraction and optimizer updates for eager sessions, including the CUDA paths that keep parameter mirrors and moment buffers on device.

def Runtime.Autograd.Torch.Internal.EagerSession.backwardDenseAllCuda {α : Type} [TorchLean.Storage α] [TensorTransfer α] (s : EagerSession α) {sh : Spec.Shape} (out : TensorRef α sh) (seed : TorchLean.Tensor α sh) : IO (Array LibTorch.AnyBuffer)

Run reverse-mode backprop on the CUDA tape, returning device gradients for all tape entries.

This is the CUDA analogue of backwardDenseAll, but it does not download gradients back to the host. This is primarily useful for implementing GPU-native optimizer steps.

def Runtime.Autograd.Torch.Internal.EagerSession.backwardScalarParamGradsCuda {α : Type} [TorchLean.Storage α] [TensorTransfer α] (s : EagerSession α) [One α] (loss : TensorRef α Spec.Shape.scalar) (seed : α := 1) : IO CudaGradMap

Run scalar-loss CUDA backprop and return gradients only for trainable parameter leaves.

seed scales the loss cotangent. Batch training uses 1 / batch.size so it can accumulate the mean directly, without first forming a larger unnormalized sum.

The traversal is the engine's sparse CUDA sweep LibTorch.Tape.backwardSparse, which walks node ids in reverse, runs a node's VJP only when that node has received a cotangent, retires every activation cotangent as soon as it has been propagated, and keeps owned device buffers only for the ids selected by retain. Here retain selects exactly the trainable parameter leaves, so the returned map has one entry per trainable parameter and optimizer updates stay on device without paying hash-table costs at every intermediate node. A trainable parameter that receives no gradient is an error, as before.

def Runtime.Autograd.Torch.Internal.EagerSession.accumulateCudaGradMap (accumulator : IO.Ref CudaGradMap) (sample : CudaGradMap) : IO Unit

Add borrowed sample gradients into an owned accumulator without downloading them.

The accumulator owns every inserted buffer. Its caller must release it on success and failure. Parameter leaf ids are stable across recordings by scalarTrainer.

def Runtime.Autograd.Torch.Internal.EagerSession.backwardDenseAll {α : Type} [TorchLean.Storage α] [TensorTransfer α] (s : EagerSession α) [Add α] [Zero α] {sh : Spec.Shape} (out : TensorRef α sh) (seed : TorchLean.Tensor α sh) : IO (Array (Spec.SomeTensor α))

Run reverse-mode backprop and return a dense gradient array for all tape entries.

seed is the upstream gradient for out (like PyTorch's backward(gradient=...)).

def Runtime.Autograd.Torch.Internal.EagerSession.grad {α : Type} [TorchLean.Storage α] {sh : Spec.Shape} (grads : Array (Spec.SomeTensor α)) (x : TensorRef α sh) : IO (TorchLean.Tensor α sh)

Extract the gradient for a particular TensorRef from a dense gradient array.

def Runtime.Autograd.Torch.Internal.EagerSession.sgdStepAll {α : Type} [TorchLean.Storage α] [TensorTransfer α] (s : EagerSession α) [Sub α] [Mul α] [Add α] [Zero α] (lr : α) (grads : Array (Spec.SomeTensor α)) : IO Unit

Apply an SGD update to all parameters recorded via use.

PyTorch comparison: for p in params: p.data -= lr * p.grad.

def Runtime.Autograd.Torch.Internal.EagerSession.checkCudaOptimizerInputs {α : Type} [TorchLean.Storage α] (operation : String) (tape : LibTorch.Tape) (params : Std.HashMap ℕ (AnyParam α)) (grads : CudaGradMap) (state : Option CudaAdamState := none) : IO Unit

Check every trainable parameter, gradient, tape value, and optional Adam state before an optimizer changes a parameter. This prevents a malformed gradient map or checkpoint state from producing a partially applied update.

def Runtime.Autograd.Torch.Internal.EagerSession.checkCudaLearningRate (operation : String) (learningRate : Float) : IO Unit

Reject a non-finite or negative learning rate before an optimizer mutates device state.

def Runtime.Autograd.Torch.Internal.EagerSession.checkSharedCudaAdamState {α : Type} [TorchLean.Storage α] (groups : Array (ParameterGroup α)) (state : CudaAdamState) : IO Unit

A shared parameter has one optimizer history, keyed by its first recorded leaf.

If an existing state contains separate histories for later aliases, there is no justified way to merge their moments. Reject that state before updating instead of silently discarding a history. Ordinary trainer slots and their checkpoint keys are unchanged.

def Runtime.Autograd.Torch.Internal.EagerSession.sgdStepAllCudaMap {α : Type} [TorchLean.Storage α] [TensorTransfer α] (s : EagerSession α) (lr : α) (grads : CudaGradMap) : IO Unit

Apply SGD from a sparse CUDA gradient map.

This is the path used by the CUDA trainer. It updates only parameter leaves and avoids allocating zero gradients for every forward activation in the tape.

def Runtime.Autograd.Torch.Internal.EagerSession.applyCudaAdamUpdate {α : Type} [TorchLean.Storage α] (s : EagerSession α) (operation : String) (configRef : IO.Ref (Option CudaAdamConfig)) (stateRef : IO.Ref CudaAdamState) (config : CudaAdamConfig) (lrF : Float) (decayStep : Unit → Float) (grads : CudaGradMap) : IO Unit

Run the shared Adam update after each entry point has converted and validated its options.

The decay coefficient stays delayed until the kernel call: AdamW evaluates -(lr * weightDecay) for each parameter after the moment corrections, while Adam passes the literal zero coefficient.

def Runtime.Autograd.Torch.Internal.EagerSession.adamStepAllCudaMap {α : Type} [TorchLean.Storage α] [TensorTransfer α] (s : EagerSession α) (configRef : IO.Ref (Option CudaAdamConfig)) (stateRef : IO.Ref CudaAdamState) (lr beta1 beta2 epsilon : α) (grads : CudaGradMap) : IO Unit

Apply Adam using an already-computed sparse CUDA gradient map.

def Runtime.Autograd.Torch.Internal.EagerSession.adamWStepAllCudaMap {α : Type} [TorchLean.Storage α] [TensorTransfer α] (s : EagerSession α) (configRef : IO.Ref (Option CudaAdamConfig)) (stateRef : IO.Ref CudaAdamState) (lr weightDecay beta1 beta2 epsilon : α) (grads : CudaGradMap) : IO Unit

Apply AdamW from a sparse CUDA gradient map.

Normal training uses this sparse map so activation gradients can be released as soon as their contributions have been propagated.