Torch Trainer Helpers

Ops instances and scalar trainer construction for the Torch-style runtime. Shape-indexed mutable parameter storage lives in Trainer.Parameters.

@[reducible, inline]

abbrev Runtime.Autograd.Torch.Internal.EagerM (α : Type) : Type → Type

Monad used for the eager Ops instance: read an Internal.EagerSession α and execute in IO.

This is the backend that makes Ops programs execute immediately by mutating a hidden runtime tape.

@[implicit_reducible]

instance Runtime.Autograd.Torch.instOpsEagerMOfTensorConv {α : Type} [Context α] [Internal.CudaBridge.TensorConv α] [DecidableEq Spec.Shape] : Ops (Internal.EagerM α) α

Ops instance for the eager Torch-style runtime.

This interprets Ops primitives by immediately executing them against the hidden mutable tape in the current Internal.EagerSession.

@[implicit_reducible]

instance Runtime.Autograd.Torch.instOpsMWith {α Δ : Type} [Context α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} : Ops (Compiled.GraphM.MWith α Δ Γ) α

Ops instance for the compiled graph-building monad GraphM.

This interprets Ops primitives by recording typed IR nodes (rather than executing immediately). See Runtime.Autograd.Compiled.GraphM and Torch.LinkedSession for how these graphs are later run.

structure Runtime.Autograd.Torch.OptimizerStateCheckpoint : Type

Persistence hooks for optimizer state owned by a backend-specific trainer.

The payload is intentionally opaque to callers. CUDA eager training, for example, owns Adam's moment buffers on the device and streams them without first constructing host tensors. A trainer that has no hidden optimizer state leaves this hook absent.

• save : System.FilePath → IO Unit

  Save the complete backend-owned optimizer state.

• load : System.FilePath → IO Unit

  Replace the backend-owned optimizer state from a previously saved payload.

structure Runtime.Autograd.Torch.ScalarTrainer (α : Type) (paramShapes inputShapes : List Spec.Shape) (natInputShapes : List Spec.Shape := []) : Type

Bundle a scalar-loss training loop for a fixed parameter pack and input signature.

This is the low-level trainer object used by module-backed execution:

• forward computes a scalar loss,
• lossAndBackward computes that loss and its parameter gradients from one tape,
• backward exposes just the gradients when the loss is not needed,
• stepWithLoss applies an SGD update and returns the loss from the same tape,
• step applies the update without requiring callers to read the loss,
• getParams reads current parameter values.

• params : ParamList α paramShapes

  Mutable trainable parameter pack.

• forward : Curried.Fn α inputShapes (Curried.Fn ℕ natInputShapes (IO (Spec.Tensor α Spec.Shape.scalar)))

  Compute the scalar loss for a curried input pack.

• lossAndBackward : Curried.Fn α inputShapes (Curried.Fn ℕ natInputShapes (IO (Spec.Tensor α Spec.Shape.scalar × TList α paramShapes)))

  Compute the scalar loss and parameter gradients from one forward tape.

• backward : Curried.Fn α inputShapes (Curried.Fn ℕ natInputShapes (IO (TList α paramShapes)))

  Compute gradients aligned with paramShapes for a curried input pack.

• stepWithLoss : α → Curried.Fn α inputShapes (Curried.Fn ℕ natInputShapes (IO (Spec.Tensor α Spec.Shape.scalar)))

  Apply one SGD-style update and return the loss used to compute that update.

• step : α → Curried.Fn α inputShapes (Curried.Fn ℕ natInputShapes (IO Unit))

  Apply one SGD-style update for a curried input pack.

• adamStep? : Option (α → α → α → α → Curried.Fn α inputShapes (Curried.Fn ℕ natInputShapes (IO Unit)))

  Optional Adam update path.

  In eager CUDA mode this is a device-gradient/device-moment update path. Other backends expose none and should use the generic optimizer wrappers.

• adamStepWithLoss? : Option (α → α → α → α → Curried.Fn α inputShapes (Curried.Fn ℕ natInputShapes (IO (Spec.Tensor α Spec.Shape.scalar))))

  CUDA-native Adam update that also returns the loss from its forward tape.

• adamWStep? : Option (α → α → α → α → α → Curried.Fn α inputShapes (Curried.Fn ℕ natInputShapes (IO Unit)))

  Optional AdamW update path.

  In eager CUDA mode this is a device-gradient/device-moment update path with decoupled weight decay. Other backends expose none and should use the generic optimizer wrappers.

• adamWStepWithLoss? : Option (α → α → α → α → α → Curried.Fn α inputShapes (Curried.Fn ℕ natInputShapes (IO (Spec.Tensor α Spec.Shape.scalar))))

  CUDA-native AdamW update that also returns the loss from its forward tape.

• optimizerStateCheckpoint? : Option OptimizerStateCheckpoint

  Save and restore optimizer state retained inside the selected runtime backend.

• getParams : IO (TList α paramShapes)

  Read current parameter values, synchronizing device mirrors if needed.

def Runtime.Autograd.Torch.Internal.gradsOfRefs {α : Type} [DecidableEq Spec.Shape] {ss : List Spec.Shape} : Array (AnyTensor α) → RefList (TensorRef α) ss → IO (TList α ss)

Extract gradients (as a typed TList) for a list of eager TensorRefs from a dense gradient array.

def Runtime.Autograd.Torch.Internal.useParams {α : Type} [CudaBridge.TensorConv α] [DecidableEq Spec.Shape] {ss : List Spec.Shape} : ParamList α ss → EagerM α (RefList (TensorRef α) ss)

Record all parameters as tape leaves in an eager session, returning their corresponding TensorRefs.

This is the eager analogue of "using" a parameter pack during a forward pass.

def Runtime.Autograd.Torch.Internal.useInputs {α : Type} [CudaBridge.TensorConv α] [DecidableEq Spec.Shape] {ss : List Spec.Shape} : TList α ss → EagerM α (RefList (TensorRef α) ss)

Record all input tensors as tape leaves in an eager session, returning their corresponding TensorRefs.

def Runtime.Autograd.Torch.scalarTrainer {α : Type} [Context α] [Internal.CudaBridge.TensorConv α] [DecidableEq Spec.Shape] {paramShapes inputShapes natInputShapes : List Spec.Shape} (opts : Options := { }) (initRequiresGrad : List Bool := List.replicate paramShapes.length true) (loss : {m : Type → Type} → [Monad m] → [inst : Ops m α] → CurriedRef (fun (s : Spec.Shape) => Ops.Ref m α s) (paramShapes ++ inputShapes) (CurriedRef (fun (s : Spec.Shape) => Ops.NatTensorRef m α s) natInputShapes (m (Ops.Ref m α Spec.Shape.scalar)))) : Curried.Fn α paramShapes (IO (ScalarTrainer α paramShapes inputShapes natInputShapes))

Build a ScalarTrainer from an initial parameter pack and a backend-generic loss definition.

loss is written once against the Ops interface over a concatenated context paramShapes ++ inputShapes. Depending on opts.backend, we either:

• compile the loss once (compiled backend), or
• execute it eagerly by building a runtime tape each step (eager backend).

def Runtime.Autograd.Torch.scalarTrainer.gradsPrefix {α : Type} [DecidableEq Spec.Shape] {ss : List Spec.Shape} : Array (AnyTensor α) → ℕ → IO (TList α ss)