Module Execution

This file connects typed module definitions to executable scalar modules. It provides parameter initialization, execution settings, and helpers that select a scalar type and device from command line arguments.

ScalarModuleDef describes the forward and loss programs together with their parameter and input shapes. Instantiating it produces a mutable ScalarModule that can evaluate inputs, run backward, and update parameters. The shape lists remain part of both types, so construction and execution use the same parameter ordering.

def TorchLean.Module.instantiateConfigured {α : Type} [Context α] [DecidableEq Spec.Shape] [Runtime.Autograd.Torch.Internal.CudaBridge.TensorConv α] {paramShapes inputShapes natInputShapes : List Spec.Shape} (defn : ScalarModuleDef paramShapes inputShapes natInputShapes) (cast : Float → α) (opts : Options) : IO (ScalarModule α paramShapes inputShapes natInputShapes)

Instantiate a ScalarModuleDef under explicit Torch options such as backend and device.

The supplied options are passed unchanged to module construction, including the selected device and execution strategy.

def TorchLean.Module.instantiateFloatWithPlan {paramShapes inputShapes natInputShapes : List Spec.Shape} (defn : ScalarModuleDef paramShapes inputShapes natInputShapes) (opts : Options) (plan : RuntimeInit.Plan paramShapes) : IO (ScalarModule Float paramShapes inputShapes natInputShapes)

Instantiate a Float module with runtime layer parameter initializers.

The initializer plan is indexed by the same paramShapes list as the module, so Lean checks that every parameter has exactly one initializer. In CUDA mode, supported initializers allocate device buffers directly instead of first constructing every parameter as a large nested Lean tensor.

def TorchLean.Module.instantiateFloatWithInit {paramShapes inputShapes natInputShapes : List Spec.Shape} (defn : ScalarModuleDef paramShapes inputShapes natInputShapes) (opts : Options) (inits : List RuntimeInit.FloatInit) : IO (ScalarModule Float paramShapes inputShapes natInputShapes)

List-based wrapper for checkpoint/JSON boundaries.

If the caller has a statically known parameter list, prefer instantiateFloatWithPlan; this wrapper checks the list length before applying it.

structure TorchLean.Module.ExecConfig : Type

Execution configuration parsed from CLI flags.

Supported flags (parsed by ExecConfig.parseAndStrip):

• --dtype ... / --float32-mode ... (see TorchLean.Runtime.DType)
• --backend eager|compiled
• --device auto|cpu|cuda|rocm|metal|wasm|tpu|trainium|custom|external
• --show-backend (print backend capsules when the eager runtime first executes them)

• dtype : Runtime.DType

  Scalar dtype selection.

• backend : Backend

  Execution backend selection.

• device : NN.Backend.Device

  Explicit eager execution device.

• showBackend : Bool

  Print each backend capsule when the eager runtime first executes it.

def TorchLean.Module.instReprExecConfig.repr : ExecConfig → ℕ → Std.Format

@[implicit_reducible]

instance TorchLean.Module.instReprExecConfig : Repr ExecConfig

@[implicit_reducible]

instance TorchLean.Module.instDecidableEqExecConfig : DecidableEq ExecConfig

def TorchLean.Module.instDecidableEqExecConfig.decEq (x✝ x✝¹ : ExecConfig) : Decidable (x✝ = x✝¹)

def TorchLean.Module.ExecConfig.parseBackend (v : String) : Except String Backend

Parse a backend selector string into a runtime Backend.

def TorchLean.Module.ExecConfig.parseDevice (value : String) : Except String NN.Backend.Device

Parse a CLI device selector. auto currently resolves to the portable CPU runtime.

def TorchLean.Module.ExecConfig.requestsCuda : List String → Bool

Whether a raw CLI argument list explicitly requests CUDA.

def TorchLean.Module.ExecConfig.parseAndStripWithDefaultDType (args : List String) (defaultDType : Runtime.DType) : Except String (ExecConfig × List String)

Parse CLI flags handled by ExecConfig and return (cfg, rest).

Consumed flags:

• --backend eager|compiled (at most once),
• --device auto|cpu|cuda|rocm|metal|wasm|tpu|trainium|custom|external,
• --show-backend (boolean flag; removed from rest).

All dtype/Float32 selection flags are delegated to TorchLean.Runtime.DType.parseAndStripWithDefault.

Default dtype policy:

• If the user does not specify --dtype / --float32-mode and CUDA is selected, default to dtype=float (CUDA eager supports Float upload/download).
• Otherwise default to dtype=float32 (executable IEEE-754 float32 semantics).

Named future devices are accepted at parse time so --show-backend and planning diagnostics can explain them. Runtime session creation still rejects devices that this build cannot execute.

The selected device chooses its normal registered kernels. Users do not need a second performance flag after selecting CUDA or another accelerator.

def TorchLean.Module.ExecConfig.parseAndStripWithDefaultDType.go (device : NN.Backend.Device) (showBackend : Bool) (acc : List String) : List String → Except String (NN.Backend.Device × Bool × List String)

def TorchLean.Module.ExecConfig.toOptions (cfg : ExecConfig) (seed : ℕ := 0) : Except String Options

Convert a parsed CLI execution config to runtime Options.

def TorchLean.Module.ExecConfig.parseAndStrip (args : List String) : Except String (ExecConfig × List String)

Parse CLI flags with the standard TorchLean default dtype policy.

def TorchLean.Module.ExecConfig.log (cfg : ExecConfig) : IO Unit

Log the chosen execution config to stdout for reproducible runs.

def TorchLean.Module.withRuntime (args : List String) (k : {α : Type} → [Context α] → [DecidableEq Spec.Shape] → [ToString α] → [Runtime.FromFloat α] → (Float → α) → Options → List String → IO Unit) : IO Unit

Parse runtime flags (--dtype, --backend, --device, --show-backend) and choose an executable scalar α, then call k with:

• cast : Float → α for building inputs from literals
• opts : Options selecting the backend/kernel mode
• rest : List String containing the remaining CLI arguments

This is useful for scripts that need to build a dataset/loader (and maybe determine shapes/batch sizes) before instantiating a concrete ScalarModuleDef.

def TorchLean.Module.withModule {paramShapes inputShapes : List Spec.Shape} (defn : ScalarModuleDef paramShapes inputShapes) (args : List String) (k : {α : Type} → [inst : Context α] → [inst_1 : DecidableEq Spec.Shape] → [ToString α] → (Float → α) → ScalarModule α paramShapes inputShapes → List String → IO Unit) : IO Unit

Instantiate a ScalarModuleDef under CLI runtime flags (--dtype, --backend, --device, --show-backend), then call a continuation.

This provides the cast function Float → α so call sites can build inputs from float literals.

def TorchLean.Module.withModuleRuntime {paramShapes inputShapes : List Spec.Shape} (defn : ScalarModuleDef paramShapes inputShapes) (args : List String) (k : {α : Type} → [inst : Context α] → [inst_1 : DecidableEq Spec.Shape] → [ToString α] → [Runtime.FromFloat α] → ScalarModule α paramShapes inputShapes → List String → IO Unit) : IO Unit

Like withModule, but also provides an _root_.TorchLean.Runtime.FromFloat α instance (for numeric literals).