Torch Utils

Helpers for writing PyTorch-style training loops on top of Runtime.Autograd.Torch.

This file focuses on training-loop ergonomics:

• extract scalar values,
• build short TLists,
• run simple SGD loops for Torch.ScalarTrainer.

Stateful optimizer loops live in Runtime.Autograd.TorchLean, because those depend on TorchLean.Optim. Keeping that dependency out of this low-level utility module prevents the session/ref layer from depending upward on the model/optimizer API surface.

Initialization helpers (Float constants)

inductive Runtime.Autograd.Torch.Init.Scheme : Type

Deterministic initialization schemes stored as Float constants.

PyTorch comparison:

• .zeros / .ones correspond to torch.nn.init.zeros_ / torch.nn.init.ones_
• .uniform lo hi corresponds to torch.nn.init.uniform_ (with explicit a=lo, b=hi)
• .xavierUniform fanIn fanOut corresponds to torch.nn.init.xavier_uniform_
• .kaimingUniform fanIn corresponds to torch.nn.init.kaiming_uniform_ with nonlinearity="relu"

References:

• https://pytorch.org/docs/stable/nn.init.html
• https://pytorch.org/docs/stable/generated/torch.nn.init.xavier_uniform_.html
• https://pytorch.org/docs/stable/generated/torch.nn.init.kaiming_uniform_.html

• zeros : Scheme
• ones : Scheme
• uniform (lo hi : Float) : Scheme
• normal (mean std : Float) : Scheme
• xavierUniform (fanIn fanOut : ℕ) : Scheme
• kaimingUniform (fanIn : ℕ) : Scheme

@[implicit_reducible]

instance Runtime.Autograd.Torch.Init.instReprScheme : Repr Scheme

def Runtime.Autograd.Torch.Init.instReprScheme.repr : Scheme → ℕ → Std.Format

def Runtime.Autograd.Torch.Init.Internal.splitmix64 (x : UInt64) : UInt64

SplitMix64 mixing used by indexed parameter initialization.

def Runtime.Autograd.Torch.Init.Internal.rand01 (seed idx : ℕ) : Float

Deterministic U[0,1) sampler derived independently from a seed and scalar index.

Indexing is constant-time, so constructing an n-element tensor takes O(n) sampler work. The same key/index formula is used by TorchLean's storage-first runtime initializer.

def Runtime.Autograd.Torch.Init.Internal.sampleAt (sch : Scheme) (seed idx : ℕ) : Float

Sample the idx-th scalar of a tensor initialized using Scheme.

This is the scalar-level primitive used by Init.tensor.

def Runtime.Autograd.Torch.Init.tensor (sch : Scheme) (seed : ℕ := 0) {s : Spec.Shape} : Spec.Tensor Float s

Create a Tensor Float s by sampling a Scheme deterministically.

This pure initializer is convenient for model definitions and reproducible examples. Runtime initialization paths can provide more specialized allocation strategies for very large tensors.

PyTorch comparison: this mimics using torch.nn.init.* routines on freshly allocated parameters, but here we work with pure Tensor Float s values (no mutation) and use a deterministic seeded sampler for reproducibility.

def Runtime.Autograd.Torch.Init.tensor.build (sch : Scheme) (seed' idxBase : ℕ) {t : Spec.Shape} : ℕ → Spec.Tensor Float t

def Runtime.Autograd.Torch.Init.xavierW (outDim inDim : ℕ) (seed : ℕ := 0) : Spec.Tensor Float (Spec.Shape.dim outDim (Spec.Shape.dim inDim Spec.Shape.scalar))

Xavier/Glorot-uniform initializer for 2D weight matrices.

PyTorch comparison: torch.nn.init.xavier_uniform_ with gain=1.

def Runtime.Autograd.Torch.Init.kaimingW (outDim inDim : ℕ) (seed : ℕ := 0) : Spec.Tensor Float (Spec.Shape.dim outDim (Spec.Shape.dim inDim Spec.Shape.scalar))

Kaiming/He-uniform initializer for 2D weight matrices.

PyTorch comparison: torch.nn.init.kaiming_uniform_ with nonlinearity="relu" and default parameters (so the bound is sqrt(6/fan_in)).

Conveniences for scalar training loops

@[reducible, inline]

abbrev Runtime.Autograd.Torch.scalarOf {α : Type} (t : Spec.Tensor α Spec.Shape.scalar) : α

Extract the scalar value from a scalar-shaped tensor.

PyTorch comparison: like t.item() for a 0-dim tensor.

def Runtime.Autograd.Torch.tlistSingleton {α : Type} {s₁ : Spec.Shape} (x₁ : Spec.Tensor α s₁) : TList α [s₁]

Build a one-element TList (useful for curried trainer APIs).

TList syntax sugar

def Runtime.Autograd.Torch.tlistBang : Lean.ParserDescr

Build a TList from a comma-separated list of terms.

This is meant for training code where tlistSingleton/tlistPair/… becomes tedious.

Example:

let xs : TList Float [.dim 2 .scalar, .dim 1 .scalar] :=
  tlist![x, y]

def Runtime.Autograd.Torch.tlistPair {α : Type} {s₁ s₂ : Spec.Shape} (x₁ : Spec.Tensor α s₁) (x₂ : Spec.Tensor α s₂) : TList α [s₁, s₂]

Build a two-element TList (useful for curried trainer APIs).

def Runtime.Autograd.Torch.tlistTriple {α : Type} {s₁ s₂ s₃ : Spec.Shape} (x₁ : Spec.Tensor α s₁) (x₂ : Spec.Tensor α s₂) (x₃ : Spec.Tensor α s₃) : TList α [s₁, s₂, s₃]

Build a three-element TList (useful for curried trainer APIs).

def Runtime.Autograd.Torch.tlistQuad {α : Type} {s₁ s₂ s₃ s₄ : Spec.Shape} (x₁ : Spec.Tensor α s₁) (x₂ : Spec.Tensor α s₂) (x₃ : Spec.Tensor α s₃) (x₄ : Spec.Tensor α s₄) : TList α [s₁, s₂, s₃, s₄]

Build a four-element TList (useful for curried trainer APIs).

def Runtime.Autograd.Torch.ScalarTrainer.forwardT {α : Type} {paramShapes inputShapes natInputShapes : List Spec.Shape} (tr : ScalarTrainer α paramShapes inputShapes natInputShapes) (xs : TList α inputShapes) (natInputs : TList ℕ natInputShapes) : IO (Spec.Tensor α Spec.Shape.scalar)

Uncurried forward pass for ScalarTrainer.

ScalarTrainer.forward is stored as a curried function over the input shapes; this helper lets you pass a TList (like a tuple of tensors).

def Runtime.Autograd.Torch.ScalarTrainer.lossAndBackwardT {α : Type} {paramShapes inputShapes natInputShapes : List Spec.Shape} (tr : ScalarTrainer α paramShapes inputShapes natInputShapes) (xs : TList α inputShapes) (natInputs : TList ℕ natInputShapes) : IO (Spec.Tensor α Spec.Shape.scalar × TList α paramShapes)

Uncurried loss-and-gradient pass for ScalarTrainer.

The loss and gradients come from the same tape. Use this instead of calling forwardT followed by backwardT when both results are needed.

def Runtime.Autograd.Torch.ScalarTrainer.backwardT {α : Type} {paramShapes inputShapes natInputShapes : List Spec.Shape} (tr : ScalarTrainer α paramShapes inputShapes natInputShapes) (xs : TList α inputShapes) (natInputs : TList ℕ natInputShapes) : IO (TList α paramShapes)

Uncurried backward pass for ScalarTrainer.

Returns per-parameter gradients (aligned with paramShapes).

def Runtime.Autograd.Torch.ScalarTrainer.stepT {α : Type} {paramShapes inputShapes natInputShapes : List Spec.Shape} (tr : ScalarTrainer α paramShapes inputShapes natInputShapes) (lr : α) (xs : TList α inputShapes) (natInputs : TList ℕ natInputShapes) : IO Unit

Uncurried SGD step for ScalarTrainer.

PyTorch comparison: analogous to loss.backward(); optimizer.step() for a fixed SGD optimizer, except here the trainer bundles the update rule.

def Runtime.Autograd.Torch.ScalarTrainer.stepWithLossT {α : Type} {paramShapes inputShapes natInputShapes : List Spec.Shape} (tr : ScalarTrainer α paramShapes inputShapes natInputShapes) (lr : α) (xs : TList α inputShapes) (natInputs : TList ℕ natInputShapes) : IO (Spec.Tensor α Spec.Shape.scalar)

Uncurried SGD step that returns the loss used to compute the update.

def Runtime.Autograd.Torch.trainCycleSGD {α : Type} [ToString α] {paramShapes inputShapes : List Spec.Shape} (tr : ScalarTrainer α paramShapes inputShapes) (lr : α) (steps : ℕ) (samples : List (TList α inputShapes)) (logEvery : ℕ := 1) : IO Unit

Train steps SGD updates, cycling through samples.

PyTorch comparison: this matches the common eager training skeleton:

for step in range(steps):
  batch = dataset[step % len(dataset)]
  loss = forward(batch)
  step(lr, batch)   # typically: loss.backward(); optimizer.step()

Note: ScalarTrainer.step is the "bundled SGD optimizer" for the trainer. Stateful optimizers (Adam, RMSProp, ...) are exposed from Runtime.Autograd.TorchLean.

def Runtime.Autograd.Torch.meanLoss {α : Type} [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] {paramShapes inputShapes : List Spec.Shape} (tr : ScalarTrainer α paramShapes inputShapes) (samples : List (TList α inputShapes)) : IO α

Evaluate mean loss over a dataset.

PyTorch comparison: like running a model in torch.no_grad() over a dataloader and averaging the scalar loss values, except here we call ScalarTrainer.forwardT directly.