Runtime Initialization

Casting helpers and shape-indexed initialization plans for executable modules. Runtime initializers can materialize parameter storage on the host or directly in CUDA buffers.

Small helpers

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

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

This is mainly used to turn ordinary Float tensor literals into Float/ExecFloat.Binary 8 23/etc.

def Runtime.Autograd.Model.Module.castPack {α : Type} [TorchLean.Storage α] (cast : Float → α) {ss : List Spec.Shape} : TorchLean.TensorPack Float ss → TorchLean.TensorPack α ss

List-shaped castTensor for TorchLean's TorchLean.TensorPack parameter bundles.

Runtime Float Initializers

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

Runtime initializer for a Float parameter.

The usual ObjectiveDef.initState 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.

• normal (mean std : Float) (seed : ℕ := 0) : FloatInit

  Normal distribution with explicit mean and standard deviation.

• 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.

def Runtime.Autograd.Model.Module.RuntimeInit.FloatInit.ofScheme (scheme : Torch.Init.Scheme) (seed : ℕ := 0) : FloatInit

Translate a proof-visible initializer scheme into its storage-first runtime form.

def Runtime.Autograd.Model.Module.RuntimeInit.FloatInit.validate (expectedSize : ℕ) : FloatInit → Except String Unit

Validate one runtime initializer before it allocates or mutates parameter storage.

inductive Runtime.Autograd.Model.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 TorchLean.TensorPack: the type says there is exactly one initializer for each parameter shape, in the same order. That removes the 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. Proofs still concern 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.Model.Module.RuntimeInit.Plan.append {ss₁ ss₂ : List Spec.Shape} : Plan ss₁ → Plan ss₂ → Plan (ss₁ ++ ss₂)

Concatenate two shape-indexed initializer plans.

def Runtime.Autograd.Model.Module.RuntimeInit.Plan.toArray {ss : List Spec.Shape} : Plan ss → Array FloatInit

Forget the shape index when interoperating with runtime-sized callers.

theorem Runtime.Autograd.Model.Module.RuntimeInit.Plan.size_toArray {ss : List Spec.Shape} (plan : Plan ss) : plan.toArray.size = ss.length

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

def Runtime.Autograd.Model.Module.RuntimeInit.Plan.validate {ss : List Spec.Shape} : Plan ss → Except String Unit

Validate every initializer against its parameter shape before applying the plan.

def Runtime.Autograd.Model.Module.RuntimeInit.Plan.ofArray? (ss : List Spec.Shape) (inits : Array FloatInit) : Except String (Plan ss)

Recover a shape-indexed plan from a runtime-sized initializer array.

def Runtime.Autograd.Model.Module.RuntimeInit.Plan.ofArray?.go (inits : Array FloatInit) (remaining : List Spec.Shape) : ℕ → Except String (Plan remaining)

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

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

def Runtime.Autograd.Model.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:

$$ \begin{aligned} \operatorname{fanIn} &=\operatorname{inChannels}\,k_1\cdots k_d,\\ \operatorname{fanOut} &=\operatorname{outChannels}\,k_1\cdots k_d. \end{aligned} $$

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

def Runtime.Autograd.Model.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.Model.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.Model.Module.RuntimeInit.unitAt (seed idx : ℕ) : Float

Deterministic unit sample shared with the pure tensor initializer.

Calling the canonical sampler here keeps CPU storage-first initialization equal to the semantic tensor initializer. The CUDA path uses the same SplitMix64 key/index sequence, materialized as float32 device values.

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

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

def Runtime.Autograd.Model.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.Model.Module.RuntimeInit.natToU32Checked (ctx : String) (n : ℕ) : IO UInt32

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

def Runtime.Autograd.Model.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.Model.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.Model.Module.RuntimeInit.hostTensorOf {α : Type} [TorchLean.Storage α] (cast : Float → α) {s : Spec.Shape} (init : FloatInit) : IO (TorchLean.Tensor α s)

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

def Runtime.Autograd.Model.Module.RuntimeInit.zeroPack {α : Type} [TorchLean.Storage α] (zero : α) {ss : List Spec.Shape} : TorchLean.TensorPack α 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.Model.Module.RuntimeInit.Internal.applyPlanUnchecked {α : Type} [TorchLean.Storage α] [Torch.TensorTransfer α] (cast : Float → α) (options : Config) {ss : List Spec.Shape} : Torch.ParamList α ss → Plan ss → IO Unit

Apply a plan after the public entrypoint has validated every initializer.

def Runtime.Autograd.Model.Module.RuntimeInit.applyPlan {α : Type} [TorchLean.Storage α] [Torch.TensorTransfer α] (cast : Float → α) (options : Config) {ss : List Spec.Shape} (parameters : Torch.ParamList α ss) (plan : Plan ss) : IO Unit

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

The shape list appears on both sides of the type:

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

The complete plan is validated before any parameter is mutated, so an invalid later initializer cannot leave the module partially initialized.