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 α. The public API (Trainer.Session, Trainer.train, Trainer.predict) wraps it and exposes Float only.

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), reusable no-gradient forward and evaluation-loss evaluators over the same live state, and the current mode. 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) [σ, τ]
• modeRef : IO.Ref TorchLean.nn.Mode

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) [σ, τ]) (modeRef : IO.Ref nn.Mode) : Runner α model

Construct a runner from its objective, evaluators, and mode cell.

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 one execution mode.

opaque TorchLean.Trainer.Internal.Runner.evaluationEvaluator {σ τ : 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.

opaque TorchLean.Trainer.Internal.Runner.modeCell {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) : IO.Ref nn.Mode

The runner's mode cell.

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 α] : IO (Runner α model)

Instantiate a model and objective under a runtime scalar, injecting literals 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.readFloatPack {α : Type} [Storage α] [Context α] [Runtime.Autograd.Torch.TensorTransfer α] {shapes : List Shape} : TensorPack α shapes → IO (TensorPack Float shapes)

Read every tensor of a runtime pack back to host Float.

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

Read the state back as host Float tensors; exact for binary32 runtime scalars.

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

Replace the state from host Float tensors, casting them into the runtime scalar.

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.mode {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) : IO nn.Mode

Read the runner's current mode (.train or .eval).

def TorchLean.Trainer.Internal.Runner.setMode {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) (value : nn.Mode) : IO Unit

Set the runner mode (.train or .eval).

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

Select training behavior for stateful layers.

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

Select evaluation behavior for stateful layers.

def TorchLean.Trainer.Internal.Runner.forwardWithMode {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) (selectedMode : nn.Mode) (input : Tensor α σ) : IO (Tensor α τ)

Evaluate one input tensor in an explicit mode without changing the runner's mode cell.

def TorchLean.Trainer.Internal.Runner.forward {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) (input : Tensor α σ) : IO (Tensor α τ)

Evaluate one input tensor using the active mode (.train or .eval).

def TorchLean.Trainer.Internal.Runner.predict {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) (input : Tensor α σ) : IO (Tensor α τ)

Run evaluation-mode prediction without changing the runner's persistent mode.

def TorchLean.Trainer.Internal.Runner.sampleLossWithMode {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) (selectedMode : nn.Mode) (sample : Sample.Supervised α σ τ) : IO α

Scalar loss of one supervised sample in an explicit mode without changing the runner's mode cell.

Training uses the instantiated objective, including its random stream and buffer updates. Evaluation uses a reusable no-gradient evaluator over the same live parameter objects and does not update running buffers.

def TorchLean.Trainer.Internal.Runner.sampleLoss {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) (sample : Sample.Supervised α σ τ) : IO α

Scalar loss of one supervised sample using the active mode.

def TorchLean.Trainer.Internal.Runner.meanLoss {σ τ : Shape} {model : nn.Sequential σ τ} {α : Type} [Storage α] [Context α] (runner : Runner α model) (samples : Data.SampleStream (Sample.Supervised α σ τ)) : IO α

Mean scalar loss over a finite sample stream in the active mode; 0 for an empty stream.

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.

def TorchLean.Trainer.Internal.withBoundOptimizer.mapStateList {α : Type} [Storage α] {State : (α : Type) → [Storage α] → Shape → Type} {shapes : List Shape} : ({s : Shape} → State α s → State α s) → Runtime.Autograd.Model.StateList State α shapes → Runtime.Autograd.Model.StateList State α shapes

Gradient accumulation and optimizer updates

def TorchLean.Trainer.Internal.addGradients {α : Type} [Storage α] [Add α] {shapes : List Shape} (first second : nn.State α shapes) : nn.State α shapes

Add two shape-aligned model-state gradients.

def TorchLean.Trainer.Internal.scaleGradients {α : Type} [Storage α] [Mul α] {shapes : List Shape} (factor : α) (gradient : nn.State α shapes) : nn.State α shapes

Scale every tensor in a model-state gradient.

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.stepBatch {σ τ : 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 σ τ} {α : Type} [Storage α] [Context α] (stepper : Stepper α model) (sample : Sample.Supervised α σ τ) : IO α

Run one optimizer step and return its scalar loss.

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

Run one averaged-gradient optimizer step over a nonempty batch and return its mean scalar loss.

Every sample is differentiated at the same parameter point. The completed-step counter advances once for the whole batch.

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

Run one optimizer step over a nonempty batch without 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 switches the runner to training mode, 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.