Torch Utils

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

This file focuses on demo/training ergonomics:

• extract scalar values,
• build small 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 facade.

Initialization helpers (Float constants)

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

Deterministic initialization schemes for small demos (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
• xavierUniform (fanIn fanOut : ℕ) : Scheme
• kaimingUniform (fanIn : ℕ) : Scheme

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

@[implicit_reducible]

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

def Runtime.Autograd.Torch.Init.Internal.lcgNext (s : ℕ) : ℕ

Next state of a simple 32-bit linear congruential generator (LCG).

This exists to make demos deterministic and reproducible; it is not statistically strong and it is not cryptographically secure.

def Runtime.Autograd.Torch.Init.Internal.lcg01 (s : ℕ) : Float

Interpret the 32-bit LCG state as a Float in [0, 1).

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

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

Implementation note: this is the LCG iterated idx+1 times, then normalized to [0,1).

def Runtime.Autograd.Torch.Init.Internal.rand01.go : ℕ → ℕ → ℕ

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 is intended for small demo models. It is not optimized.

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

Small Sample Generators (Float Constants)

def Runtime.Autograd.Torch.Samples.vec2 (x1 x2 : Float) : Spec.Tensor Float (Spec.Shape.dim 2 Spec.Shape.scalar)

Turn a point (x1,x2) into a Tensor Float (.dim 2 .scalar).

def Runtime.Autograd.Torch.Samples.vec1 (y : Float) : Spec.Tensor Float (Spec.Shape.dim 1 Spec.Shape.scalar)

Turn a scalar y into a Tensor Float (.dim 1 .scalar).

def Runtime.Autograd.Torch.Samples.affine2 (w1 w2 b x1 x2 : Float) : Float

Affine map y = w1*x1 + w2*x2 + b for building small regression datasets.

Small conveniences for scalar training demos

@[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.tlist1 {α : 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 small demo/training code, where tlist1/tlist2/… becomes tedious.

Example:

def xs : TList Float [Shape.Vec 2, Shape.Vec 1] :=
  tlist![x, y]

def Runtime.Autograd.Torch.tlist2 {α : 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.tlist3 {α : 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.tlist4 {α : 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 : List Spec.Shape} (tr : ScalarTrainer α paramShapes inputShapes) (xs : TList α inputShapes) : 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.backwardT {α : Type} {paramShapes inputShapes : List Spec.Shape} (tr : ScalarTrainer α paramShapes inputShapes) (xs : TList α inputShapes) : IO (TList α paramShapes)

Uncurried backward pass for ScalarTrainer.

Returns per-parameter gradients (aligned with paramShapes).

def Runtime.Autograd.Torch.ScalarTrainer.stepT {α : Type} {paramShapes inputShapes : List Spec.Shape} (tr : ScalarTrainer α paramShapes inputShapes) (lr : α) (xs : TList α inputShapes) : 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.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.