Module

TorchLean module wrappers with PyTorch-style ergonomics.

TorchLean already provides the core ingredients:

• a small Ops interface, so you write a model once and run it on different backends;
• scalarTrainer, which builds an eager or compiled training loop for scalar losses.

This file adds a thin “nn.Module-style” wrapper so users can:

• package initial parameters + a loss definition as a single object,
• instantiate it under a chosen backend (.eager / .compiled),
• call forward / backward / step / params with a small, consistent API.

Important: dtype selection is handled in NN.API.DType (because it picks the Lean type α). The module definitions here are polymorphic in α, so the same module can be:

• used in executables with Float / IEEE32Exec, or
• instantiated at ℝ in proofs (noncomputable; not for IO execution).

Small helpers

def Runtime.Autograd.TorchLean.Module.castTensor {α : Type} (cast : Float → α) {s : Spec.Shape} (t : Spec.Tensor Float s) : Spec.Tensor α s

Cast a Float tensor to a backend scalar type α by mapping a scalar cast function.

This is mainly used to turn tensorND!-authored Float initializers into Float/IEEE32Exec/etc.

def Runtime.Autograd.TorchLean.Module.castTList {α : Type} (cast : Float → α) {ss : List Spec.Shape} : TList Float ss → TList α ss

List-shaped version of castTensor for TorchLean's TList parameter bundles.

Runtime Float Initializers

inductive Runtime.Autograd.TorchLean.Module.RuntimeInit.FloatInit : Type

Runtime initializer for a Float parameter.

The usual ScalarModuleDef.initParams path stores initializers as typed Lean tensors. That is the right representation when the initial value itself is part of the Lean object being inspected. For large Float runs, it is better to allocate runtime storage from a compact initialization scheme and synchronize the host tensor only when parameters are explicitly read back.

The design mirrors the storage-first APIs used by mainstream runtimes:

• PyTorch exposes in-place initializers such as torch.nn.init.uniform_, torch.nn.init.xavier_uniform_, and torch.nn.init.kaiming_uniform_ for already-allocated tensors: https://pytorch.org/docs/stable/nn.init.html.
• PyTorch's meta-device / to_empty path separates "module structure exists" from "real storage is materialized", after which users explicitly initialize parameters: https://docs.pytorch.org/docs/main/meta.html.

TorchLean keeps the semantic parameter type (Tensor Float s) available, but this runtime path lets CPU/CUDA execution initialize real storage directly.

• zeros : FloatInit

  Fill with zeros. PyTorch analogue: torch.nn.init.zeros_.

• ones : FloatInit

  Fill with ones. PyTorch analogue: torch.nn.init.ones_.

• uniform (lo hi : Float) (seed : ℕ := 0) : FloatInit

  Uniform distribution over [lo, hi), using TorchLean's deterministic runtime RNG.

• xavierUniform (fanIn fanOut : ℕ) (seed : ℕ := 0) : FloatInit

  Xavier/Glorot uniform with explicit fan-in and fan-out.

• kaimingUniform (fanIn : ℕ) (seed : ℕ := 0) : FloatInit

  Kaiming/He uniform with explicit fan-in.

• flat (values : FloatArray) : FloatInit

  Exact row-major payload. Used for imported checkpoints or generated tensors.

inductive Runtime.Autograd.TorchLean.Module.RuntimeInit.Plan : List Spec.Shape → Type

A shape-indexed initialization plan.

This is the typed runtime-initialization API for modules with a known parameter shape list. It is the initialization analogue of TList: the type says there is exactly one initializer for each parameter shape, in the same order. That removes the annoying runtime failure mode where a plain list is one element too short or too long.

The initializers themselves are runtime schemes rather than proof objects. The proof-facing story is still the ordinary Tensor Float s parameter value; this plan only controls how the executable Float runtime materializes those tensors on CPU or CUDA.

• nil : Plan []

  No parameters, no initializers.

• cons {s : Spec.Shape} {ss : List Spec.Shape} (init : FloatInit) (rest : Plan ss) : Plan (s :: ss)

  Initializer for the head parameter, followed by the plan for the remaining parameters.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.Plan.toList {ss : List Spec.Shape} : Plan ss → List FloatInit

Forget the shape index when interoperating with list-based callers.

theorem Runtime.Autograd.TorchLean.Module.RuntimeInit.Plan.length_toList {ss : List Spec.Shape} (plan : Plan ss) : plan.toList.length = ss.length

The type index is not decorative: forgetting a Plan ss to a list produces exactly ss.length initializers. This checked fact lets the runtime API avoid the usual "initializer list does not match parameter list" class of bugs once a plan has been built.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.Plan.ofList? (ss : List Spec.Shape) : List FloatInit → Except String (Plan ss)

Recover a shape-indexed plan from a plain list.

List-based callers still enter through this boundary, but the runtime converts them immediately into the shape-indexed representation before touching any parameters.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.dimProduct (xs : List ℕ) : ℕ

Product of a list of dimensions, used for convolutional receptive-field sizes.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.fanInOut? (s : Spec.Shape) : Option (ℕ × ℕ)

Infer (fanIn, fanOut) from a parameter shape using the common linear/conv convention.

For a matrix shaped [out, in], this returns (in, out). For convolution-like weights shaped [outChannels, inChannels, k1, ..., kd], it returns:

fanIn  = inChannels  * k1 * ... * kd
fanOut = outChannels * k1 * ... * kd

This is the same fan convention documented by PyTorch's Xavier/Kaiming initialization utilities.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.xavierUniformForShape (s : Spec.Shape) (seed : ℕ := 0) : Except String FloatInit

Build a Xavier initializer by deriving fan-in/fan-out from a Linear/Conv-style weight shape.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.kaimingUniformForShape (s : Spec.Shape) (seed : ℕ := 0) : Except String FloatInit

Build a Kaiming initializer by deriving fan-in from a Linear/Conv-style weight shape.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.xavierLinearWeight (outDim inDim : ℕ) (seed : ℕ := 0) : FloatInit

Convenience initializer for a matrix weight stored as [outDim, inDim].

def Runtime.Autograd.TorchLean.Module.RuntimeInit.kaimingLinearWeight (_outDim inDim : ℕ) (seed : ℕ := 0) : FloatInit

Convenience initializer for a ReLU-style matrix weight stored as [outDim, inDim].

def Runtime.Autograd.TorchLean.Module.RuntimeInit.unitAt (seed idx : ℕ) : Float

Deterministic unit sample used by CPU/runtime initialization.

The CUDA path below uses Cuda.Buffer.randUniform, which is keyed by the same SplitMix64 family. Exact CPU/CUDA bit equality is not the contract here; reproducible initialization for a fixed path is. Tests that need exact CUDA RNG parity use the lower-level CUDA RNG stress tests.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.sampleAt : FloatInit → ℕ → Float

Scalar value generated by a FloatInit at a row-major flat index.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.floatArrayOf (n : ℕ) (init : FloatInit) : IO FloatArray

Materialize an initializer as a host FloatArray.

CPU execution uses this path directly. CUDA uses it only when the initializer already is an exact flat payload; analytic initializers such as uniform/Xavier/Kaiming are created on the runtime side.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.natToU32Checked (ctx : String) (n : ℕ) : IO UInt32

Checked conversion to the current CUDA buffer API's UInt32 element count.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.cudaUniformBuffer (n : ℕ) (lo hi : Float) (seed : ℕ) : IO Cuda.Buffer

Allocate a CUDA buffer filled with U(lo, hi).

The implementation keeps all element generation on the runtime side: first create a CUDA uniform buffer in [0,1), then perform lo + (hi-lo) * u with CUDA buffer ops.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.cudaBufferOf (n : ℕ) (init : FloatInit) : IO Cuda.Buffer

Allocate a CUDA buffer for a FloatInit.

For analytic schemes (zeros, ones, uniform, xavierUniform, kaimingUniform), this avoids building a large nested Lean tensor. For .flat, the caller already supplied the exact payload, so we upload that payload directly.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.hostTensorOf {s : Spec.Shape} (init : FloatInit) : IO (Spec.Tensor Float s)

Materialize a runtime initializer as a normal host tensor. Used for CPU execution.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.zeroFloatTList {ss : List Spec.Shape} : TList Float ss

Host slots for a parameter list before runtime initialization installs the real values.

CUDA runtime initialization immediately replaces these with CUDA mirrors and marks the host values stale. These entries still give the existing Param type a valid host slot for later explicit readback.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.applyFloatPlan (opts : Options) {ss : List Spec.Shape} : Torch.ParamList Float ss → Plan ss → IO Unit

Apply a shape-indexed initialization plan to an already-created parameter list.

The key point is that the shape list appears on both sides of the type:

Torch.ParamList Float ss → RuntimeInit.Plan ss → IO Unit

So Lean checks the bookkeeping that Python frameworks usually check at runtime: every parameter gets one initializer, and no extra initializer is silently ignored.

def Runtime.Autograd.TorchLean.Module.RuntimeInit.applyFloatInits (opts : Options) {ss : List Spec.Shape} (ps : Torch.ParamList Float ss) (inits : List FloatInit) : IO Unit

Apply a runtime list of initializers after checking it against the parameter shapes.

Plan ss is the typed form used by the initializer engine. This entrypoint is for places where the initializer list comes from outside Lean's typechecker, such as a checkpoint, JSON file, or CLI experiment.

Scalar-loss module (training)

structure Runtime.Autograd.TorchLean.Module.ScalarModuleDef (paramShapes inputShapes : List Spec.Shape) : Type 1

A scalar-loss module definition:

• initParams is stored as Float constants (easy to write with tensorND!),
• loss is polymorphic in the scalar backend (same code works for Float/IEEE32Exec/…).

You can instantiate this definition as a ScalarModule under a chosen backend and dtype.

• initParams : TList Float paramShapes

  Initial parameter values, stored as Float tensors and cast at instantiation time.

• initRequiresGrad : List Bool

  Per-parameter requiresGrad flags aligned with paramShapes.

• loss {α : Type} [Context α] [DecidableEq Spec.Shape] : Program α (paramShapes ++ inputShapes) Spec.Shape.scalar

  Scalar loss program over (params ++ inputs), polymorphic in the scalar backend.

structure Runtime.Autograd.TorchLean.Module.ScalarModule (α : Type) [Context α] [DecidableEq Spec.Shape] (paramShapes inputShapes : List Spec.Shape) : Type

Runtime module instance (the thing you "run").

This wraps Torch.ScalarTrainer, but exposes a more Module-like set of methods.

• trainer : Torch.ScalarTrainer α paramShapes inputShapes

def Runtime.Autograd.TorchLean.Module.ScalarModule.create {α : Type} [Context α] [DecidableEq Spec.Shape] [Torch.Internal.CudaBridge.TensorConv α] {paramShapes inputShapes : List Spec.Shape} (opts : Options := { }) (initRequiresGrad : List Bool := List.replicate paramShapes.length true) (loss : {m : Type → Type} → [Monad m] → [inst : Ops m α] → CurriedRef (fun (s : Spec.Shape) => Torch.Ops.Ref m α s) (paramShapes ++ inputShapes) (m (Torch.Ops.Ref m α Spec.Shape.scalar))) (initParams : TList α paramShapes) : IO (ScalarModule α paramShapes inputShapes)

Create a runtime scalar-loss module from an explicit loss program and initial parameter values.

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

def Runtime.Autograd.TorchLean.Module.ScalarModule.forward {α : Type} [Context α] [DecidableEq Spec.Shape] {paramShapes inputShapes : List Spec.Shape} (m : ScalarModule α paramShapes inputShapes) (xs : TList α inputShapes) : IO (Spec.Tensor α Spec.Shape.scalar)

Run the forward pass and return the scalar loss value.

def Runtime.Autograd.TorchLean.Module.ScalarModule.backward {α : Type} [Context α] [DecidableEq Spec.Shape] {paramShapes inputShapes : List Spec.Shape} (m : ScalarModule α paramShapes inputShapes) (xs : TList α inputShapes) : IO (TList α paramShapes)

Run one forward/backward pass and return gradients for all parameters.

def Runtime.Autograd.TorchLean.Module.ScalarModule.step {α : Type} [Context α] [DecidableEq Spec.Shape] {paramShapes inputShapes : List Spec.Shape} (m : ScalarModule α paramShapes inputShapes) (lr : α) (xs : TList α inputShapes) : IO Unit

Convenience "one-step SGD": compute gradients and apply an SGD update with learning rate lr.

def Runtime.Autograd.TorchLean.Module.ScalarModule.initOptim {α : Type} [Context α] [DecidableEq Spec.Shape] {paramShapes inputShapes : List Spec.Shape} (m : ScalarModule α paramShapes inputShapes) (opt : Optim.Optimizer α paramShapes) : IO opt.State

Initialize an optimizer state for this module's parameters.

def Runtime.Autograd.TorchLean.Module.ScalarModule.stepWith {α : Type} [Context α] [DecidableEq Spec.Shape] {paramShapes inputShapes : List Spec.Shape} (m : ScalarModule α paramShapes inputShapes) (opt : Optim.Optimizer α paramShapes) (st : opt.State) (xs : TList α inputShapes) : IO opt.State

Run one optimizer step using an explicit optimizer + state.

This mirrors a PyTorch training step:

1. compute gradients (backwardT)
2. update parameters via opt.step and return the new optimizer state

def Runtime.Autograd.TorchLean.Module.ScalarModule.params {α : Type} [Context α] [DecidableEq Spec.Shape] {paramShapes inputShapes : List Spec.Shape} (m : ScalarModule α paramShapes inputShapes) : IO (TList α paramShapes)

Fetch the current parameter values as a shape-indexed list.

def Runtime.Autograd.TorchLean.Module.ScalarModule.setParams {α : Type} [Context α] [DecidableEq Spec.Shape] {paramShapes inputShapes : List Spec.Shape} (m : ScalarModule α paramShapes inputShapes) (ps : TList α paramShapes) : IO Unit

Overwrite all parameter values.

def Runtime.Autograd.TorchLean.Module.ScalarModule.trainSGD {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] {paramShapes inputShapes : List Spec.Shape} (m : ScalarModule α paramShapes inputShapes) (lr : α) (steps : ℕ) (samples : List (TList α inputShapes)) (logEvery : ℕ := 1) : IO Unit

Train with vanilla SGD for a fixed number of steps on a fixed list of samples.

def Runtime.Autograd.TorchLean.Module.ScalarModule.trainWith {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] {paramShapes inputShapes : List Spec.Shape} (m : ScalarModule α paramShapes inputShapes) (opt : Optim.Optimizer α paramShapes) (st0 : opt.State) (steps : ℕ) (samples : List (TList α inputShapes)) (logEvery : ℕ := 1) : IO opt.State

Like trainSGD, but with an explicit optimizer + mutable optimizer state.

def Runtime.Autograd.TorchLean.Module.ScalarModule.meanLoss {α : Type} [Context α] [DecidableEq Spec.Shape] [ToString α] {paramShapes inputShapes : List Spec.Shape} (m : ScalarModule α paramShapes inputShapes) (samples : List (TList α inputShapes)) : IO α

Compute the mean loss over a list of samples (no parameter updates).

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

Instantiate a ScalarModuleDef by casting Float initializers to α and choosing Torch options.

This is the most general constructor. The shorter instantiate entrypoint chooses standard runtime options before calling this function.

def Runtime.Autograd.TorchLean.Module.ScalarModuleDef.instantiateFloatWithRuntimePlan {paramShapes inputShapes : List Spec.Shape} (d : ScalarModuleDef paramShapes inputShapes) (opts : Options) (plan : RuntimeInit.Plan paramShapes) : IO (ScalarModule Float paramShapes inputShapes)

Instantiate a Float module using runtime parameter initializers.

This is the runtime-initialized sibling of instantiateWith. Instead of first building every initial parameter as a full Lean tensor, it creates minimal zero host tensors and then applies a shape-indexed runtime plan to the module parameters. In CUDA mode those initializers allocate device buffers directly and mark the host copies stale; public parameter readback still synchronizes them through the existing CUDA mirror machinery.

def Runtime.Autograd.TorchLean.Module.ScalarModuleDef.instantiateFloatWithRuntimeInit {paramShapes inputShapes : List Spec.Shape} (d : ScalarModuleDef paramShapes inputShapes) (opts : Options) (inits : List RuntimeInit.FloatInit) : IO (ScalarModule Float paramShapes inputShapes)

Instantiate a Float module from a plain initializer list.

This wrapper is useful at file/JSON boundaries. Internally it immediately checks the list against paramShapes and then delegates to instantiateFloatWithRuntimePlan, so the actual parameter mutation still goes through the shape-indexed path.

def Runtime.Autograd.TorchLean.Module.ScalarModuleDef.instantiate {α : Type} [Context α] [DecidableEq Spec.Shape] [Torch.Internal.CudaBridge.TensorConv α] {paramShapes inputShapes : List Spec.Shape} (d : ScalarModuleDef paramShapes inputShapes) (cast : Float → α) (backend : Backend := Torch.Backend.eager) : IO (ScalarModule α paramShapes inputShapes)

Convenience instantiator that chooses only the backend (.eager or .compiled).