Module Execution

This file connects typed module definitions to executable scalar modules. It provides state initialization and explicit execution settings. Command-line selection lives in NN.API.Module.Command.

ObjectiveDefinition describes a scalar objective together with model state and input shapes. Instantiating it produces a mutable Objective that can evaluate the objective, return explicit state gradients, and update trainable entries. The shape lists remain part of both types, so construction and execution use the same state ordering.

@[reducible, inline]

abbrev TorchLean.Module.ObjectiveDefinition (β : Type) [Storage β] (stateShapes inputShapes : List Shape) (dataInputShapes : List Shape := []) : Type 1

An immutable scalar objective together with its typed initial state.

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

Executable state for a model and scalar objective.

The runtime trainer and device machinery are intentionally hidden. Public code interacts with an objective through operations such as loss, grad, step, and state.

• runtimeObjective : Runtime.Autograd.Model.Module.Objective α β stateShapes inputShapes dataInputShapes

@[reducible, inline]

abbrev TorchLean.Module.Evaluator (α β : Type) [Storage α] [Storage β] [Context α] (stateShapes inputShapes dataInputShapes : List Shape) (output : Shape) : Type

A reusable evaluator over an existing typed model state.

opaque TorchLean.Module.Objective.Internal.fromRuntime {α β : Type} [Storage α] [Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Shape} (objective : Runtime.Autograd.Model.Module.Objective α β stateShapes inputShapes dataInputShapes) : Objective α β stateShapes inputShapes dataInputShapes

Wrap a runtime objective behind the public execution boundary.

opaque TorchLean.Module.Objective.Internal.runtime {α β : Type} [Storage α] [Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Shape} (objective : Objective α β stateShapes inputShapes dataInputShapes) : Runtime.Autograd.Model.Module.Objective α β stateShapes inputShapes dataInputShapes

Reveal the runtime objective only to implementation code.

def TorchLean.Module.Objective.loss {α β : Type} [Storage α] [Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Shape} (module : Objective α β stateShapes inputShapes dataInputShapes) (inputs : Arguments α inputShapes) (dataInputs : Arguments β dataInputShapes) : IO (Tensor α [])

Evaluate the scalar objective.

def TorchLean.Module.Objective.grad {α β : Type} [Storage α] [Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Shape} (module : Objective α β stateShapes inputShapes dataInputShapes) (inputs : Arguments α inputShapes) (dataInputs : Arguments β dataInputShapes) (value : Bool := false) : IO (match value with | false => nn.State α stateShapes | true => nn.State α stateShapes × Tensor α [])

Differentiate the objective with respect to its state.

Entries that do not require gradients are zero. Set value := true to return (gradient, objectiveValue) from one forward tape.

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

Start a fresh optimizer history from the module's current parameters and buffers.

This clears backend-owned moments and permits switching between native and generic updates. Use the returned optimizer state for subsequent steps.

def TorchLean.Module.Objective.step {α β : Type} [Storage α] [Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Shape} (module : Objective α β stateShapes inputShapes dataInputShapes) (optimizer : Runtime.Autograd.Model.Optim.Optimizer α stateShapes) (optimizerState : optimizer.State) (inputs : Arguments α inputShapes) (dataInputs : Arguments β dataInputShapes) (loss : Bool := false) : IO (match loss with | false => optimizer.State | true => optimizer.State × Tensor α [])

Compute one gradient update.

By default this returns the next optimizer state. Set loss := true to return (nextOptimizerState, lossValue) from the same training tape.

def TorchLean.Module.Objective.Internal.tryNativeBatchStep {α β : Type} [Storage α] [Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Shape} (module : Objective α β stateShapes inputShapes dataInputShapes) (optimizer : Runtime.Autograd.Model.Optim.Optimizer α stateShapes) (optimizerState : optimizer.State) (batch : Array (Arguments α inputShapes × Arguments β dataInputShapes)) (readLoss : Bool) : IO (Option (optimizer.State × Option (Tensor α [])))

Try one backend-native mean-gradient update over a nonempty batch.

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

Apply a generic optimizer update from an already-computed, shape-aligned state gradient.

On CUDA, this cannot share an optimizer history with native steps: their device moments are separate from the explicit optimizer state. Call initOptimizer before switching paths.

def TorchLean.Module.Objective.state {α β : Type} [Storage α] [Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Shape} (module : Objective α β stateShapes inputShapes dataInputShapes) : IO (nn.State α stateShapes)

Read the complete parameter-and-buffer state as host tensors.

On CUDA, parameters changed on the device are copied back before this function returns. The runtime retains these host values and reuses them while they remain current. Reading the state after another device update requires a fresh transfer and host storage for the updated parameters.

Use this operation when inspecting, saving, or transferring a model's state. For large models, include the host copy in the memory needed for those operations.

def TorchLean.Module.Objective.setState {α β : Type} [Storage α] [Storage β] [Context α] {stateShapes inputShapes dataInputShapes : List Shape} (module : Objective α β stateShapes inputShapes dataInputShapes) (state : nn.State α stateShapes) : IO Unit

Replace the complete parameter-and-buffer state.

def TorchLean.Module.Internal.instantiate {α β : Type} [Storage α] [Storage β] [Context α] [Runtime.Autograd.Torch.TensorTransfer α] {stateShapes inputShapes dataInputShapes : List Shape} (definition : ObjectiveDefinition β stateShapes inputShapes dataInputShapes) (cast : Float → α) (runtime : Runtime.Config := { }) (initialState? : Option (nn.State α stateShapes) := none) : IO (Objective α β stateShapes inputShapes dataInputShapes)

Instantiate an objective with an explicit scalar-literal conversion.

def TorchLean.Module.instantiate {β : Type} [Storage β] {stateShapes inputShapes dataInputShapes : List Shape} (definition : ObjectiveDefinition β stateShapes inputShapes dataInputShapes) (runtime : Runtime.Config := { }) (α : Type := Float32) [Storage α] [Context α] [Runtime.FromFloat α] [Runtime.Autograd.Torch.TensorTransfer α] (initialState? : Option (nn.State α stateShapes) := none) : IO (Objective α β stateShapes inputShapes dataInputShapes)

Instantiate an executable objective, optionally starting from exact typed state.

initialState? bypasses the model's stored Float initializers. Without it, initialization uses the runtime arithmetic's standard Float conversion.

This is the low-level constructor for custom losses and multi-input programs. The higher-level nn.Module and Trainer APIs should be preferred for ordinary sequential models.