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} [CudaBridge.TensorConv α] (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (out : TensorRef α sh) (seed : Spec.Tensor α sh) : IO (Array Cuda.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.backwardScalarDenseAllCuda {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Add α] [Zero α] [One α] [DecidableEq Spec.Shape] (loss : TensorRef α Spec.Shape.scalar) : IO (Array Cuda.AnyBuffer)

Run CUDA backward from a scalar loss with seed 1, returning device gradient buffers.

def Runtime.Autograd.Torch.Internal.EagerSession.backwardScalarParamGradsCuda {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [One α] [DecidableEq Spec.Shape] (loss : TensorRef α Spec.Shape.scalar) : IO CudaGradMap

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

The tape walk uses an array indexed by node id rather than a persistent hash map. Only trainable parameter gradients are packed into the returned map, so optimizer updates stay on device without paying hash-table costs at every intermediate node.

def Runtime.Autograd.Torch.Internal.EagerSession.backwardDenseAll {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Add α] [Zero α] [DecidableEq Spec.Shape] {sh : Spec.Shape} (out : TensorRef α sh) (seed : Spec.Tensor α sh) : IO (Array (AnyTensor α))

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.backwardScalarDenseAll {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Add α] [Zero α] [One α] [DecidableEq Spec.Shape] (loss : TensorRef α Spec.Shape.scalar) : IO (Array (AnyTensor α))

Run backward from a scalar loss with seed 1.

PyTorch comparison: loss.backward() for a scalar loss.

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

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

def Runtime.Autograd.Torch.Internal.EagerSession.sgdStepAll {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [Sub α] [Mul α] [Add α] [Zero α] [DecidableEq Spec.Shape] (lr : α) (grads : Array (AnyTensor α)) : 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.sgdStepAllCuda {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [DecidableEq Spec.Shape] (lr : α) (grads : Array Cuda.AnyBuffer) : IO Unit

Apply an SGD update to all parameters recorded via use, using CUDA device gradients.

This avoids downloading the full dense gradient array and keeps updated parameters in each Param's CUDA mirror. Host tensors are synchronized later by explicit parameter readback.

def Runtime.Autograd.Torch.Internal.EagerSession.checkCudaOptimizerInputs {α : Type} (operation : String) (tape : Cuda.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.sgdStepAllCudaMap {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [DecidableEq Spec.Shape] (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.adamStepAllCudaMap {α : Type} [CudaBridge.TensorConv α] (s : EagerSession α) [DecidableEq Spec.Shape] (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} [CudaBridge.TensorConv α] (s : EagerSession α) [DecidableEq Spec.Shape] (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.