Scalar Objectives

Executable scalar-objective definitions and runtime state. This module provides loss and gradient evaluation, explicit optimizer steps, state access, and optimizers bound to a live objective.

structure Runtime.Autograd.Model.Module.ObjectiveDef (β : Type) [TorchLean.Storage β] (stateShapes inputShapes : List Spec.Shape) (dataInputShapes : List Spec.Shape := []) : Type 1

An immutable scalar-objective definition:

• initState stores initial trainable parameters and persistent buffers as Float tensors,
• loss is polymorphic in the scalar backend (same code works for Float/configured binary32/…).

You can instantiate this definition as an Objective under a chosen execution mode and scalar.

• initState : TorchLean.TensorPack Float stateShapes

  Initial parameter-and-buffer state, cast from Float at instantiation time.

• runtimeInit : Option (RuntimeInit.Plan stateShapes)

  Optional storage-first initialization plan for executable Float runs.

  The ordinary tensors remain the semantic initial values. This plan records how a runtime may materialize the same initialization directly in backend storage without traversing a large nested Lean tensor first.

• requiresGrad : Array Bool

  Differentiability flags aligned with stateShapes; persistent buffers carry false.

• validate : Except String Unit

  Validate static model and objective configuration before runtime allocation.

• validateDataInputs : TorchLean.TensorPack β dataInputShapes → Except String Unit

  Validate non-differentiable inputs before they reach the runtime program.

• loss {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] : CurriedRef (fun (s : Spec.Shape) => Torch.Ops.Ref m α s) (stateShapes ++ inputShapes) (CurriedRef (fun (s : Spec.Shape) => Torch.Ops.DataRef m α β s) dataInputShapes (m (Torch.Ops.Ref m α Spec.Shape.scalar)))

  Scalar loss over differentiable tensors followed by non-differentiable data tensors.

  The second curried input pack can carry labels, bounded token IDs, or gather indices without converting them through the model's floating-point scalar type.

structure Runtime.Autograd.Model.Module.Objective (α β : Type) [TorchLean.Storage α] [TorchLean.Storage β] [Context α] (stateShapes inputShapes : List Spec.Shape) (dataInputShapes : List Spec.Shape := []) : Type

Runtime state for a model together with a scalar objective.

This is lower level than PyTorch's loss classes: it owns model state as well as the objective. It wraps Torch.ScalarTrainer and exposes objective evaluation, explicit gradients, and updates.

• trainer : Torch.ScalarTrainer α β stateShapes inputShapes dataInputShapes

  Trainer that owns trainable parameters and persistent buffers.

• runtime : Config

  Runtime configuration used to instantiate the module.

• tensorTransfer : Torch.TensorTransfer α

  Concrete host/device tensor conversion selected when the module was instantiated.

def Runtime.Autograd.Model.Module.Objective.create {α β : Type} [TorchLean.Storage α] [TorchLean.Storage β] [Context α] [Torch.TensorTransfer α] {stateShapes inputShapes dataInputShapes : List Spec.Shape} (runtime : Config := { }) (requiresGrad : Array Bool := Array.replicate stateShapes.length true) (validateDataInputs : TorchLean.TensorPack β dataInputShapes → Except String Unit := fun (x : TorchLean.TensorPack β dataInputShapes) => pure ()) (loss : {m : Type → Type} → [Monad m] → [inst : Ops m α] → CurriedRef (fun (s : Spec.Shape) => Torch.Ops.Ref m α s) (stateShapes ++ inputShapes) (CurriedRef (fun (s : Spec.Shape) => Torch.Ops.DataRef m α β s) dataInputShapes (m (Torch.Ops.Ref m α Spec.Shape.scalar)))) (initState : TorchLean.TensorPack α stateShapes) : IO (Objective α β stateShapes inputShapes dataInputShapes)

Create a runtime objective from an explicit scalar program and initial model state.

This is the low-level constructor; public training code starts from an ObjectiveDef and calls ObjectiveDef.instantiate.

def Runtime.Autograd.Model.Module.Objective.loss {α β : Type} [TorchLean.Storage α] [TorchLean.Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Spec.Shape} (m : Objective α β stateShapes inputShapes dataInputShapes) (xs : TorchLean.TensorPack α inputShapes) (dataInputs : TorchLean.TensorPack β dataInputShapes) : IO (TorchLean.Tensor α Spec.Shape.scalar)

Evaluate the scalar objective.

def Runtime.Autograd.Model.Module.Objective.grad {α β : Type} [TorchLean.Storage α] [TorchLean.Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Spec.Shape} (m : Objective α β stateShapes inputShapes dataInputShapes) (xs : TorchLean.TensorPack α inputShapes) (dataInputs : TorchLean.TensorPack β dataInputShapes) (value : Bool := false) : IO (match value with | false => TorchLean.TensorPack α stateShapes | true => TorchLean.TensorPack α stateShapes × TorchLean.Tensor α [])

Return state-shaped gradients, using zero for entries that do not require gradients.

Set value := true to return (gradient, objectiveValue) from one forward tape.

def Runtime.Autograd.Model.Module.Objective.initOptimizer {α β : Type} [TorchLean.Storage α] [TorchLean.Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Spec.Shape} (m : Objective α β stateShapes inputShapes dataInputShapes) (opt : Optim.Optimizer α stateShapes) : IO opt.State

Start a fresh optimizer history from this module's current state.

This also releases backend-owned moments and clears the selected update path. Use the returned state for subsequent updates; reinitialization permits switching between generic and native steps.

def Runtime.Autograd.Model.Module.Objective.step {α β : Type} [TorchLean.Storage α] [TorchLean.Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Spec.Shape} (m : Objective α β stateShapes inputShapes dataInputShapes) (opt : Optim.Optimizer α stateShapes) (st : opt.State) (xs : TorchLean.TensorPack α inputShapes) (dataInputs : TorchLean.TensorPack β dataInputShapes) (loss : Bool := false) : IO (match loss with | false => opt.State | true => opt.State × TorchLean.Tensor α [])

Run one optimizer step using an explicit optimizer and state.

This mirrors a PyTorch training step:

1. compute the explicit state gradient (ScalarTrainer.runGrad)
2. update parameters via opt.step and return the new optimizer state

Set loss := true to return (nextOptimizerState, lossValue) from the same training tape. On CUDA, mixing native steps and generic updates within one optimizer history is rejected. Call initOptimizer to start a fresh history before switching paths.

def Runtime.Autograd.Model.Module.Objective.state {α β : Type} [TorchLean.Storage α] [TorchLean.Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Spec.Shape} (m : Objective α β stateShapes inputShapes dataInputShapes) : IO (TorchLean.TensorPack α stateShapes)

Read the complete parameter-and-buffer state as a shape-indexed list.

def Runtime.Autograd.Model.Module.Objective.loadState {α β : Type} [TorchLean.Storage α] [TorchLean.Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Spec.Shape} (m : Objective α β stateShapes inputShapes dataInputShapes) (ps : TorchLean.TensorPack α stateShapes) : IO Unit

Replace the complete parameter-and-buffer state.

structure Runtime.Autograd.Model.Module.BoundOptimizer (α β : Type) [TorchLean.Storage α] [TorchLean.Storage β] [Context α] (stateShapes inputShapes dataInputShapes : List Spec.Shape) (State : Type) : Type

Mutable optimizer state bound to one executable module.

• module : Objective α β stateShapes inputShapes dataInputShapes
• state : IO.Ref State
• step : TorchLean.TensorPack α inputShapes → TorchLean.TensorPack β dataInputShapes → IO Unit

def Runtime.Autograd.Model.Module.bindOptimizer {α β : Type} [TorchLean.Storage α] [TorchLean.Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Spec.Shape} (module : Objective α β stateShapes inputShapes dataInputShapes) (optimizer : Optim.Optimizer α stateShapes) : IO (BoundOptimizer α β stateShapes inputShapes dataInputShapes optimizer.State)

Initialize an optimizer and bind its state and update operation to module.