Trainer Runtime Internals

Scalar-generic implementation of the trainer: an instantiated model with reusable forward and loss evaluators (Runner), gradient accumulation and optimizer updates, and the stateful Stepper that applies a configured optimizer and schedule.

Everything here is indexed by the runtime scalar α. Trainer.openTyped preserves that scalar at the session boundary; the runtime-selected Trainer.open, Trainer.train, and Trainer.predict paths expose Float.

structure TorchLean.Trainer.Internal.Runner (α : Type) [Storage α] [Context α] {σ τ : Shape} (model : nn.Sequential σ τ) : Type

A checked model instantiated under one runtime scalar.

This bundles the imperative runtime objective (parameters and buffers stored in refs) and reusable no-gradient forward and evaluation-loss evaluators over the same live state. Training forwards update their buffers; evaluation forwards leave them unchanged.

• runtimeObjective : TorchLean.Module.Objective α Unit (Runtime.Autograd.Model.Layers.Seq.stateShapes model) [σ, τ]
• trainingPredictor : Runtime.Autograd.Model.Module.Evaluator α Unit (Runtime.Autograd.Model.Layers.Seq.stateShapes model) [σ] [] τ
• evaluationPredictor : Runtime.Autograd.Model.Module.Evaluator α Unit (Runtime.Autograd.Model.Layers.Seq.stateShapes model) [σ] [] τ
• evaluationLossEvaluator : Runtime.Autograd.Model.Module.ObjectiveEvaluator α Unit (Runtime.Autograd.Model.Layers.Seq.stateShapes model) [σ, τ]

opaque TorchLean.Trainer.Internal.Runner.create {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runtimeObjective : Module.Objective α Unit (Runtime.Autograd.Model.Layers.Seq.stateShapes model) [σ, τ]) (trainingPredictor evaluationPredictor : Runtime.Autograd.Model.Module.Evaluator α Unit (Runtime.Autograd.Model.Layers.Seq.stateShapes model) [σ] [] τ) (evaluationLossEvaluator : Runtime.Autograd.Model.Module.ObjectiveEvaluator α Unit (Runtime.Autograd.Model.Layers.Seq.stateShapes model) [σ, τ]) : Runner α model

Construct a runner from its objective and evaluators.

opaque TorchLean.Trainer.Internal.Runner.objectiveModule {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) : Module.Objective α Unit (Runtime.Autograd.Model.Layers.Seq.stateShapes model) [σ, τ]

The executable runtime module behind a runner.

opaque TorchLean.Trainer.Internal.Runner.predictor {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) (mode : nn.Mode) : Runtime.Autograd.Model.Module.Evaluator α Unit (Runtime.Autograd.Model.Layers.Seq.stateShapes model) [σ] [] τ

The reusable no-gradient evaluator for training or evaluation behavior.

opaque TorchLean.Trainer.Internal.Runner.lossEvaluator {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) : Runtime.Autograd.Model.Module.ObjectiveEvaluator α Unit (Runtime.Autograd.Model.Layers.Seq.stateShapes model) [σ, τ]

Reusable evaluator for the evaluation-mode loss.

def TorchLean.Trainer.Internal.Runner.runtime {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) : Runtime.Config

Runtime configuration used by the runner's objective.

def TorchLean.Trainer.Internal.Runner.fromRuntimeObjective {σ τ : Shape} (model : nn.Sequential σ τ) (objective : Objective τ) {α : Type} [Storage α] [Context α] [Runtime.Autograd.Torch.TensorTransfer α] (runtimeObjective : Module.Objective α Unit (Runtime.Autograd.Model.Layers.Seq.stateShapes model) [σ, τ]) : IO (Runner α model)

Finish runner construction once its executable objective has been instantiated.

def TorchLean.Trainer.Internal.Runner.instantiate {σ τ : Shape} (model : nn.Sequential σ τ) (objective : Objective τ) (options : Runtime.Autograd.Torch.Config := { }) (α : Type := Float32) [Storage α] [Context α] [Runtime.FromFloat α] [Runtime.Autograd.Torch.TensorTransfer α] (initialState? : Option (nn.State α (Runtime.Autograd.Model.Layers.Seq.stateShapes model)) := none) : IO (Runner α model)

Instantiate a model and objective under a runtime scalar.

Explicit state is retained in α; otherwise stored Float initializers are cast with ofFloat.

def TorchLean.Trainer.Internal.Runner.state {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) : IO (nn.State α (Runtime.Autograd.Model.Layers.Seq.stateShapes model))

Read the complete parameter-and-buffer state.

def TorchLean.Trainer.Internal.Runner.setState {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) (state : nn.State α (Runtime.Autograd.Model.Layers.Seq.stateShapes model)) : IO Unit

Replace the complete parameter-and-buffer state.

def TorchLean.Trainer.Internal.Runner.initOptimizer {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) (optimizer : Runtime.Autograd.Model.Optimizer α (Runtime.Autograd.Model.Layers.Seq.stateShapes model)) : IO optimizer.State

Initialize the state owned by a runtime optimizer for this runner.

def TorchLean.Trainer.Internal.Runner.forward {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) (input : Tensor α σ) (mode : nn.Mode := Runtime.Autograd.Model.Layers.Mode.train) : IO (Tensor α τ)

Evaluate one input in the given mode (training by default).

def TorchLean.Trainer.Internal.Runner.loss {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) (sample : Sample.Supervised α σ τ) (mode : nn.Mode := Runtime.Autograd.Model.Layers.Mode.train) : IO α

Scalar loss of one sample in the given mode (training by default).

Training uses the objective's random stream and buffer updates. Evaluation uses a reusable no-gradient evaluator over the live parameters without updating running buffers.

Optimizer binding

def TorchLean.Trainer.Internal.withBoundOptimizer {σ τ : Shape} {model : nn.Sequential σ τ} {α β : Type} [Storage α] [Context α] [Runtime.FromFloat α] (runner : Runner α model) (config : optim.Optimizer) (scheduler : Option Scheduler.Config) (continuation : (optimizer : Runtime.Autograd.Model.Optimizer α (Runtime.Autograd.Model.Layers.Seq.stateShapes model)) → optimizer.State → (ℕ → optimizer.State → optimizer.State) → IO β) : IO β

Bind a configured optimizer, its initialized state, and its scheduled state updater.

Gradient accumulation and optimizer updates

def TorchLean.Trainer.Internal.meanGradAndLoss {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) (batch : Array (Sample.Supervised α σ τ)) (withLoss : Bool) : IO (nn.State α (Runtime.Autograd.Model.Layers.Seq.stateShapes model) × α)

Mean parameter gradient and mean loss for a nonempty batch at one parameter point.

With withLoss := false the losses are not read back and the second component is 0. For IEEE scalars we retain a pre-weighted sum alongside the ordinary sum and use it only if the ordinary sum becomes nonfinite. Finite means are clamped to the observed componentwise range. This avoids overflow in examples such as two equal large gradients; mixed-sign cancellation can still lose precision, and this is not an accuracy guarantee.

def TorchLean.Trainer.Internal.meanGrad {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) (batch : Array (Sample.Supervised α σ τ)) (value : Bool := false) : IO (match value with | false => nn.State α (Runtime.Autograd.Model.Layers.Seq.stateShapes model) | true => nn.State α (Runtime.Autograd.Model.Layers.Seq.stateShapes model) × α)

Compute mean parameter gradients for a nonempty batch at one parameter point.

Set value := true to return (meanGradient, meanLoss).

def TorchLean.Trainer.Internal.step {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) (optimizer : Runtime.Autograd.Model.Optimizer α (Runtime.Autograd.Model.Layers.Seq.stateShapes model)) (state : optimizer.State) (useNative : Bool) (batch : Array (Sample.Supervised α σ τ)) (loss : Bool := false) : IO (match loss with | false => optimizer.State | true => optimizer.State × α)

Apply one optimizer update to a nonempty batch.

When useNative is set, supported CUDA optimizers accumulate gradients on device and use the same moment state for every batch size. Other optimizers average explicit per-sample gradients. Set loss := true to return (nextOptimizerState, meanLoss).

Stateful steppers

structure TorchLean.Trainer.Internal.Stepper (α : Type) [Storage α] [Context α] {σ τ : Shape} (model : nn.Sequential σ τ) : Type

Stateful optimizer step functions and completed-step counter for one runner.

• runBatch : Array (TorchLean.Sample.Supervised α σ τ) → IO α
• runBatchSilently : Array (TorchLean.Sample.Supervised α σ τ) → IO Unit
• stepRef : IO.Ref ℕ

opaque TorchLean.Trainer.Internal.Stepper.create {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runBatch : Array (Sample.Supervised α σ τ) → IO α) (runBatchSilently : Array (Sample.Supervised α σ τ) → IO Unit) (stepRef : IO.Ref ℕ) : Stepper α model

Construct a stepper from its hidden batch-step functions and counter.

opaque TorchLean.Trainer.Internal.Stepper.action {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (stepper : Stepper α model) : Array (Sample.Supervised α σ τ) → IO α

The hidden loss-returning batch-step function.

opaque TorchLean.Trainer.Internal.Stepper.silentAction {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (stepper : Stepper α model) : Array (Sample.Supervised α σ τ) → IO Unit

The hidden batch-step function that does not read the loss back.

opaque TorchLean.Trainer.Internal.Stepper.counter {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (stepper : Stepper α model) : IO.Ref ℕ

The hidden completed-step counter.

def TorchLean.Trainer.Internal.Stepper.step {σ τ : Shape} {model : nn.Sequential σ τ} {Input α : Type} [Storage α] [Context α] (stepper : Stepper α model) (sample : Input) (batch loss : Bool := false) [Internal.BatchInput (Sample.Supervised α σ τ) (Array (Sample.Supervised α σ τ)) batch Input] : IO (match loss with | false => Unit | true => α)

Apply one optimizer update, optionally returning its scalar loss.

batch := true takes a nonempty array and averages gradients at the same parameter point. The counter advances once for the whole batch. loss := false avoids reading the loss back.

def TorchLean.Trainer.Internal.Stepper.steps {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (stepper : Stepper α model) : IO ℕ

Read the number of completed optimizer steps.

def TorchLean.Trainer.Internal.Runner.stepper {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] [Runtime.FromFloat α] (runner : Runner α model) (optimizer : optim.Optimizer) (scheduler : Option Scheduler.Config := none) : IO (Stepper α model)

Construct a Stepper for a runner, optimizer config, and optional scheduler.

Every step trains through the runner's objective, applies the schedule for the current step index, and advances the completed-step counter once per batch. A singleton batch uses the runtime's native single-sample update.