Supervised Runtime Training

Supervised tasks, runners, steppers, optimizer configs, trainer aliases, and the low-level session exports that back executable examples.

inductive TorchLean.Trainer.Manual.SeqLoss : Type

Built-in loss choices for SeqTask.

• mse (reduction : Loss.Reduction := Loss.Reduction.mean) : SeqLoss
• crossEntropyOneHot (reduction : Loss.Reduction := Loss.Reduction.mean) : SeqLoss

structure TorchLean.Trainer.Manual.SeqTask (σ τ : Spec.Shape) : Type 2

A supervised task is just a model plus a choice of loss.

• model : Runtime.Autograd.TorchLean.NN.Seq σ τ

  Model to run.

• loss : SeqLoss

  Loss function.

def TorchLean.Trainer.Manual.SeqTask.moduleDefWithMode {σ τ : Spec.Shape} (task : SeqTask σ τ) (mode : Runtime.Autograd.TorchLean.NN.Mode) : Runtime.Autograd.TorchLean.ScalarModuleDef task.model.paramShapes [σ, τ]

Build a ScalarModuleDef for a task, choosing an explicit model mode (train/eval).

This is the underlying "instantiate me as a runnable module" step for training.

def TorchLean.Trainer.Manual.SeqTask.moduleDef {σ τ : Spec.Shape} (task : SeqTask σ τ) : Runtime.Autograd.TorchLean.ScalarModuleDef task.model.paramShapes [σ, τ]

Default module definition for a task (training mode).

def TorchLean.Trainer.Manual.SeqTask.mse {σ τ : Spec.Shape} (model : Runtime.Autograd.TorchLean.NN.Seq σ τ) (reduction : Loss.Reduction := Loss.Reduction.mean) : SeqTask σ τ

Constructor: regression task (MSE loss).

def TorchLean.Trainer.Manual.SeqTask.crossEntropyOneHot {σ τ : Spec.Shape} (model : Runtime.Autograd.TorchLean.NN.Seq σ τ) (reduction : Loss.Reduction := Loss.Reduction.mean) : SeqTask σ τ

Constructor: one-hot classification task (cross-entropy loss).

@[reducible, inline]

abbrev TorchLean.Trainer.Manual.paramShapes {σ τ : Spec.Shape} (task : SeqTask σ τ) : List Spec.Shape

Parameter shapes for a task (delegates to Seq.paramShapes).

inductive TorchLean.Trainer.Manual.OptimizerConfig : Type

Optimizer hyperparameter configuration for the supervised training helpers.

This configuration covers the optimizer choices exposed by the public training helpers. It mirrors a few common PyTorch optimizers by name/defaults, but it does not try to cover the full option surface of torch.optim.*.

• sgd (lr : Float) (momentum : Float := 0.0) : OptimizerConfig

  SGD optimizer config.

  PyTorch analogy: torch.optim.SGD(..., lr=..., momentum=...) when momentum > 0, and plain SGD when momentum = 0.

• adagrad (lr : Float) (epsilon : Float := 1e-10) : OptimizerConfig

  AdaGrad optimizer config.

• rmsprop (lr : Float) (decay : Float := 0.99) (epsilon : Float := 1e-8) : OptimizerConfig

  RMSProp optimizer config.

• adam (lr : Float) (beta1 : Float := 0.9) (beta2 : Float := 0.999) (epsilon : Float := 1e-8) : OptimizerConfig

  Adam optimizer config.

• adamw (lr : Float) (weightDecay : Float := 1e-2) (beta1 : Float := 0.9) (beta2 : Float := 0.999) (epsilon : Float := 1e-8) : OptimizerConfig

  AdamW optimizer config (decoupled weight decay).

• adadelta (lr : Float := 1.0) (rho : Float := 0.9) (epsilon : Float := 1e-6) : OptimizerConfig

  Adadelta optimizer config.

def TorchLean.Trainer.Manual.instReprOptimizerConfig.repr : OptimizerConfig → ℕ → Std.Format

@[implicit_reducible]

instance TorchLean.Trainer.Manual.instReprOptimizerConfig : Repr OptimizerConfig

structure TorchLean.Trainer.Manual.TrainConfig : Type

Step-based training configuration for trainSamples / trainDataset.

Fields:

• steps: number of parameter updates,
• batchSize: number of samples consumed by one public step for in-memory datasets,
• optimizer: optimizer hyperparameters,
• scheduler: optional learning-rate schedule (applied per step),
• logEvery: progress printing frequency (0 disables logging).

• steps : ℕ

  Number of optimizer updates.

• batchSize : ℕ

  Number of dataset items consumed by one training step.

  The loop differentiates every item at the same parameter point, averages the resulting gradient packs, and performs one optimizer update. If each item is already a fixed-size tensor minibatch, keep this value at 1 for one vectorized forward/backward pass per update.

• optimizer : OptimizerConfig

  Optimizer configuration.

• scheduler : Option Scheduler.Config

  Scheduler configuration.

• logEvery : ℕ

  Log once every this many steps.

• cudaMemWatch : ℕ

  Sample CUDA allocator state every this many completed steps; 0 disables sampling.

@[implicit_reducible]

instance TorchLean.Trainer.Manual.instReprTrainConfig : Repr TrainConfig

def TorchLean.Trainer.Manual.instReprTrainConfig.repr : TrainConfig → ℕ → Std.Format

structure TorchLean.Trainer.Manual.CudaMemWatchState : Type

State carried by the CUDA-memory drift detector used by sustained training runs.

• firstStep : ℕ
• firstFreeBytes : ℕ
• warned : Bool

def TorchLean.Trainer.Manual.instReprCudaMemWatchState.repr : CudaMemWatchState → ℕ → Std.Format

@[implicit_reducible]

instance TorchLean.Trainer.Manual.instReprCudaMemWatchState : Repr CudaMemWatchState

def TorchLean.Trainer.Manual.effectiveCudaMemWatch (opts : Runtime.Autograd.Torch.Options) (steps requested : ℕ) : ℕ

Resolve an explicit CUDA-memory cadence, or enable periodic sampling for very long runs.

def TorchLean.Trainer.Manual.reportCudaMemWatch (opts : Runtime.Autograd.Torch.Options) (watchEvery totalSteps done : ℕ) (state? : Option CudaMemWatchState) : IO (Option CudaMemWatchState)

Sample the CUDA allocator and warn when sustained free-memory loss projects exhaustion before the requested run completes.

structure TorchLean.Trainer.Manual.TrainReport (α : Type) : Type

Small summary returned by lower training helpers.

By default, before and after are mean loss values, but the type is polymorphic so callers can report other scalars in the same shape.

• before : α

  Metrics before training.

• after : α

  Metrics after training.

structure TorchLean.Trainer.Manual.LoaderTrainConfig : Type

Epoch-based training configuration for trainLoader (data-loader training).

Fields:

• epochs: number of epochs (each epoch iterates once over the loader),
• optimizer: optimizer hyperparameters,
• scheduler: optional learning-rate schedule applied once per epoch,
• logEvery: progress printing frequency (0 disables logging).

• epochs : ℕ

  Number of epochs to train for.

• optimizer : OptimizerConfig

  Optimizer configuration.

• scheduler : Option Scheduler.Config

  Scheduler configuration.

• logEvery : ℕ

  Log once every this many steps.

def TorchLean.Trainer.Manual.instReprLoaderTrainConfig.repr : LoaderTrainConfig → ℕ → Std.Format

@[implicit_reducible]

instance TorchLean.Trainer.Manual.instReprLoaderTrainConfig : Repr LoaderTrainConfig

def TorchLean.Trainer.Manual.optimizerLR : OptimizerConfig → Float

Extract the base learning rate encoded in an optimizer configuration.

def TorchLean.Trainer.Manual.stepLR (scheduler : Option Scheduler.Config) (cfg : OptimizerConfig) (step : ℕ) : Float

Resolve the learning rate to use at a given training step.

If a scheduler is present, it takes precedence over the optimizer's baked-in base learning rate. Otherwise it returns optimizerLR cfg.

def TorchLean.Trainer.Manual.mapStateList {State : Type → Spec.Shape → Type} {α : Type} {ss : List Spec.Shape} : ({s : Spec.Shape} → State α s → State α s) → Runtime.Autograd.TorchLean.StateList State α ss → Runtime.Autograd.TorchLean.StateList State α ss

Map a state update over every optimizer-state entry in a shape-indexed parameter list.

def TorchLean.Trainer.Manual.adamStateWithLR {α : Type} (lr : α) {paramShapes : List Spec.Shape} : Runtime.Autograd.TorchLean.StateList Optim.Adam.State α paramShapes → Runtime.Autograd.TorchLean.StateList Optim.Adam.State α paramShapes

Set the learning rate field of every Adam optimizer state entry to lr.

def TorchLean.Trainer.Manual.sgdStateWithLR {α : Type} (lr : α) {paramShapes : List Spec.Shape} : Runtime.Autograd.TorchLean.StateList Optim.SGD.State α paramShapes → Runtime.Autograd.TorchLean.StateList Optim.SGD.State α paramShapes

Set the learning rate field of every plain-SGD optimizer state entry to lr.

def TorchLean.Trainer.Manual.momentumSGDStateWithLR {α : Type} (lr : α) {paramShapes : List Spec.Shape} : Runtime.Autograd.TorchLean.StateList Optim.MomentumSGD.State α paramShapes → Runtime.Autograd.TorchLean.StateList Optim.MomentumSGD.State α paramShapes

Set the learning rate field of every momentum-SGD optimizer state entry to lr.

def TorchLean.Trainer.Manual.adagradStateWithLR {α : Type} (lr : α) {paramShapes : List Spec.Shape} : Runtime.Autograd.TorchLean.StateList Optim.AdaGrad.State α paramShapes → Runtime.Autograd.TorchLean.StateList Optim.AdaGrad.State α paramShapes

Set the learning rate field of every AdaGrad optimizer state entry to lr.

def TorchLean.Trainer.Manual.rmspropStateWithLR {α : Type} (lr : α) {paramShapes : List Spec.Shape} : Runtime.Autograd.TorchLean.StateList Optim.RMSProp.State α paramShapes → Runtime.Autograd.TorchLean.StateList Optim.RMSProp.State α paramShapes

Set the learning rate field of every RMSProp optimizer state entry to lr.

def TorchLean.Trainer.Manual.adamwStateWithLR {α : Type} (lr : α) {paramShapes : List Spec.Shape} : Runtime.Autograd.TorchLean.StateList Optim.AdamW.State α paramShapes → Runtime.Autograd.TorchLean.StateList Optim.AdamW.State α paramShapes

Set the learning rate field of every AdamW optimizer state entry to lr.

def TorchLean.Trainer.Manual.adadeltaStateWithLR {α : Type} (lr : α) {paramShapes : List Spec.Shape} : Runtime.Autograd.TorchLean.StateList Optim.Adadelta.State α paramShapes → Runtime.Autograd.TorchLean.StateList Optim.Adadelta.State α paramShapes

Set the learning rate field of every Adadelta optimizer state entry to lr.

structure TorchLean.Trainer.Manual.Runner (α : Type) [Context α] [DecidableEq Spec.Shape] {σ τ : Spec.Shape} (task : SeqTask σ τ) : Type

A fully instantiated supervised task runner.

This bundles:

• the imperative ScalarModule (parameters/buffers stored in refs),
• compiled forward artifacts and loss functions for both .train and .eval modes (so switching mode is low-overhead),
• and the current mode stored in an IO.Ref.

The mode influences both operator behavior (e.g. dropout/batchnorm) and whether buffers are updated during training.

• module : Runtime.Autograd.TorchLean.ScalarModule α (paramShapes task) [σ, τ]

  Instantiated scalar module storing parameters/buffers in mutable refs.

• predictorTrain : Runtime.Autograd.Torch.CompiledGraph α (paramShapes task ++ [σ]) τ

  Compiled forward predictor specialized to training-mode behavior.

• predictorEval : Runtime.Autograd.Torch.CompiledGraph α (paramShapes task ++ [σ]) τ

  Compiled forward predictor specialized to eval-mode behavior.

• lossTrain : Runtime.Autograd.Torch.CompiledScalar α (paramShapes task ++ [σ, τ])

  Compiled loss function for training-mode behavior.

• lossEval : Runtime.Autograd.Torch.CompiledScalar α (paramShapes task ++ [σ, τ])

  Compiled loss function for eval-mode behavior.

• mode : IO.Ref Runtime.Autograd.TorchLean.NN.Mode

  Mutable mode flag (.train / .eval) used by stateful layers (e.g. dropout/batchnorm).

def TorchLean.Trainer.Manual.runnerOfModule {σ τ : Spec.Shape} (task : SeqTask σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] (module : Runtime.Autograd.TorchLean.ScalarModule α (paramShapes task) [σ, τ]) : IO (Runner α task)

Finish runner construction once parameter storage has been instantiated.

def TorchLean.Trainer.Manual.instantiateConfigured {σ τ : Spec.Shape} (task : SeqTask σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] [Runtime.Autograd.Torch.Internal.CudaBridge.TensorConv α] (cast : Float → α) (opts : Runtime.Autograd.Torch.Options := { }) : IO (Runner α task)

Instantiate a Runner by explicitly providing a Float → α cast and backend.

Use this when you want to run the same task over different numeric backends (e.g. Float vs IEEE32Exec) or when you want custom literal injection.

def TorchLean.Trainer.Manual.instantiateConfiguredFloat {σ τ : Spec.Shape} (task : SeqTask σ τ) (opts : Runtime.Autograd.Torch.Options := { }) : IO (Runner Float task)

Instantiate a Float runner using storage-first parameter initialization when the model provides it. Models without a runtime plan automatically retain the ordinary tensor initializer path.

def TorchLean.Trainer.Manual.instantiateWith {σ τ : Spec.Shape} (task : SeqTask σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] [Runtime.Autograd.Torch.Internal.CudaBridge.TensorConv α] (cast : Float → α) (backend : Runtime.Autograd.Torch.Backend := Runtime.Autograd.Torch.Backend.eager) : IO (Runner α task)

Instantiate a Runner by explicitly providing a Float → α cast and a backend selector.

Instantiate a module with explicit runtime options.

def TorchLean.Trainer.Manual.instantiateWithRuntimeOptions {σ τ : Spec.Shape} (task : SeqTask σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] [Runtime.FromFloat α] [Runtime.Autograd.Torch.Internal.CudaBridge.TensorConv α] (opts : Runtime.Autograd.Torch.Options := { }) : IO (Runner α task)

Instantiate a Runner using the standard runtime literal injection _root_.TorchLean.Runtime.ofFloat.

This is the common entrypoint for executable examples.

def TorchLean.Trainer.Manual.instantiate {σ τ : Spec.Shape} (task : SeqTask σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] [Runtime.FromFloat α] [Runtime.Autograd.Torch.Internal.CudaBridge.TensorConv α] (backend : Runtime.Autograd.Torch.Backend := Runtime.Autograd.Torch.Backend.eager) : IO (Runner α task)

Instantiate a Runner using the standard runtime literal injection _root_.TorchLean.Runtime.ofFloat and a backend selector.

Instantiate a module after parsing runtime options.

def TorchLean.Trainer.Manual.run {σ τ : Spec.Shape} (task : SeqTask σ τ) (args : List String) (k : {α : Type} → [inst : Context α] → [inst_1 : DecidableEq Spec.Shape] → [ToString α] → [Runtime.FromFloat α] → Runner α task → List String → IO Unit) : IO Unit

Run a TorchLean task with CLI-style dtype/backend selection, then call k with a fully constructed runner.

This is used by lake exe entrypoints: run takes care of parsing dtype flags and instantiating the underlying module/compiled programs.

def TorchLean.Trainer.Manual.params {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO (Runtime.Autograd.Torch.TList α (paramShapes task))

Read the current parameter list from a runner.

def TorchLean.Trainer.Manual.mode {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO Runtime.Autograd.TorchLean.NN.Mode

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

def TorchLean.Trainer.Manual.setMode {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) (value : Runtime.Autograd.TorchLean.NN.Mode) : IO Unit

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

def TorchLean.Trainer.Manual.trainMode {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO Unit

Convenience: setMode runner .train.

def TorchLean.Trainer.Manual.evalMode {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO Unit

Convenience: setMode runner .eval.

def TorchLean.Trainer.Manual.isTraining {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO Bool

Predicate: are we in training mode?

def TorchLean.Trainer.Manual.activePredictor {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO (Runtime.Autograd.Torch.CompiledGraph α (paramShapes task ++ [σ]) τ)

Pick the predictor compiled for the runner's current mode (.train or .eval).

def TorchLean.Trainer.Manual.activeLoss {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO (Runtime.Autograd.Torch.CompiledScalar α (paramShapes task ++ [σ, τ]))

Pick the loss program compiled for the runner's current mode (.train or .eval).

def TorchLean.Trainer.Manual.updateRunnerBuffers {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) (sample : Runtime.Autograd.Torch.TList α [σ, τ]) : IO Unit

Refresh mode-dependent runner buffers using one supervised sample.

This mutates the module parameters only in .train mode, mirroring PyTorch-style buffer updates for layers such as normalization. In .eval mode it is a no-op.

def TorchLean.Trainer.Manual.backward {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) (sample : Runtime.Autograd.Torch.TList α [σ, τ]) : IO (Runtime.Autograd.Torch.TList α (paramShapes task))

Run one forward/backward pass on a single supervised sample and return gradients for all parameters. Unlike PyTorch's in-place .grad convention, this API returns the gradient pack explicitly.

def TorchLean.Trainer.Manual.predict {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) (x : Spec.Tensor α σ) : IO (Spec.Tensor α τ)

Predict on one input tensor using the runner's active mode (.train or .eval).

def TorchLean.Trainer.Manual.predictBatch {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) (xs : List (Spec.Tensor α σ)) : IO (List (Spec.Tensor α τ))

Predict on a list of inputs by repeatedly calling predict.

def TorchLean.Trainer.Manual.predictClass? {σ : Spec.Shape} {n : ℕ} {task : SeqTask σ (Spec.Shape.dim n Spec.Shape.scalar)} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) (x : Spec.Tensor α σ) : IO (Option (Fin n))

For classification heads: run predict, then take argmax over the logits (if defined).

def TorchLean.Trainer.Manual.accuracyOneHot {σ : Spec.Shape} {n : ℕ} {task : SeqTask σ (Spec.Shape.dim n Spec.Shape.scalar)} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) (samples : List (Runtime.Autograd.Torch.TList α [σ, Spec.Shape.dim n Spec.Shape.scalar])) : IO (ℕ × ℕ)

Compute (correct, total) for a one-hot classification dataset.

def TorchLean.Trainer.Manual.accuracyOneHot.go {σ : Spec.Shape} {n : ℕ} {task : SeqTask σ (Spec.Shape.dim n Spec.Shape.scalar)} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) (correct total : ℕ) : List (Runtime.Autograd.Torch.TList α [σ, Spec.Shape.dim n Spec.Shape.scalar]) → IO (ℕ × ℕ)

def TorchLean.Trainer.Manual.meanLoss {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (runner : Runner α task) (samples : List (Runtime.Autograd.Torch.TList α [σ, τ])) : IO α

Mean scalar loss over a list of supervised samples (uses the runner's active mode).

def TorchLean.Trainer.Manual.meanLossDataset {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] (runner : Runner α task) (dataset : Runtime.Autograd.Train.Dataset (Runtime.Autograd.Torch.TList α [σ, τ])) : IO α

Mean scalar loss over a dataset (materialized via dataset.toList).

def TorchLean.Trainer.Manual.moduleLoss {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Context α] [DecidableEq Spec.Shape] (runner : Runner α task) (sample : Runtime.Autograd.Torch.TList α [σ, τ]) : IO α

Scalar loss for one sample through the instantiated runtime module.

def TorchLean.Trainer.Manual.effectiveTrainBatchSize (n : ℕ) : ℕ

Treat 0 as the conservative single-sample step size.

def TorchLean.Trainer.Manual.nextCyclicBatch {α : Type} (context : String) (samples : List α) (restRef : IO.Ref (List α)) (batchSize : ℕ) : IO (List α)

Take the next cyclic group of samples from an in-memory training set.