TorchLean Runtime Facade

This module defines the NN.API.TorchLean namespace: the main re-export layer over the executable TorchLean runtime:

• tensor primitives (add, matmul, reshape, ...)
• derived functional ops (F.*)
• losses / norms
• autograd entry points (grad/vjp/jacobian, etc.)
• neural-network layer combinators (NN.Seq, NN.LayerDef, ...)

The goal is to expose a stable, well-namespaced "public surface" without forcing users to import internal runtime modules directly.

How This Relates To NN.API.Public

NN.API.Public is the higher-level, PyTorch-like facade intended for most examples and tutorial code. It builds on this runtime facade and adds more ergonomic constructors (record-based configs, fewer proof arguments, etc.).

PyTorch Mapping

• API.TorchLean.NN.* corresponds to pieces of torch.nn and torch.nn.functional (https://pytorch.org/docs/stable/nn.html, https://pytorch.org/docs/stable/nn.functional.html)
• API.TorchLean.Optim.* / API.TorchLean.Trainer.* correspond to torch.optim (https://pytorch.org/docs/stable/optim.html)

Most layer helpers keep the PyTorch name but append a shape/layout suffix such as CHW or NCHW so the expected tensor layout stays visible at the call site.

TorchLean differs in two crucial ways:

• Shapes are part of types, so "wrong shape" often becomes a type error.
• Some backends are proof-only; use NN.API.DType selection when writing executables.

Recommended Import

This is the canonical module for the runtime facade: import NN.API.Runtime.

For most user-facing code, prefer NN.API.Public instead. It keeps the same runtime machinery behind a smaller, more PyTorch-shaped surface, while this module stays closer to the executable TorchLean internals and lower-level compatibility re-exports.

Core Exports

Most of this namespace is a curated re-export of _root_.Runtime.Autograd.TorchLean.*, so users can import NN.API.Runtime and get a stable API surface without importing internal runtime modules.

Rough grouping:

• execution control: Backend, Options
• program interface: Ops, RefTy, Program, CompiledOut, CompiledScalar, ...
• primitive tensor ops: add, matmul, reshape, elementwise activations, pooling, ...
• training utilities: trainCycle*, meanLoss

def NN.API.TorchLean.RefList.unpack2 {Ref : Spec.Shape → Type} {s₁ s₂ : Spec.Shape} : RefList Ref [s₁, s₂] → Ref s₁ × Ref s₂

Unpack a 2-element RefList into a pair.

Optimizer Handles (PyTorch-Like)

TorchLean optimizers are purely functional in their state: opt.step returns a new state.

This small wrapper stores the optimizer state in an IO.Ref so users can write:

let h ← API.TorchLean.Optim.handle m (TorchLean.Optim.sgd lr)
h.step sample

without manually threading the optimizer state through the training loop.

structure NN.API.TorchLean.Optim.Handle (α : Type) [Context α] [DecidableEq Spec.Shape] (paramShapes inputShapes : List Spec.Shape) (State : Type) : Type

A mutable optimizer handle bound to a concrete TorchLean ScalarModule.

The internal optimizer state is stored in an IO.Ref and updated when you call h.step sample.

• module : Runtime.Autograd.TorchLean.ScalarModule α paramShapes inputShapes

  The module whose parameters will be updated in-place.

• state : IO.Ref State

  Mutable optimizer state.

• step : TList α inputShapes → IO Unit

  One training step on a single sample, updating the internal optimizer state.

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

Create an optimizer handle for a module by initializing optimizer state from the module's current parameters.

class NN.API.TorchLean.NN.AsSeqK (F : Spec.Shape → Spec.Shape → Sort u) : Sort (max 3 u)

Adapter typeclass used by the seq! macro to treat both layers and already-sequential models as composable building blocks.

This exists purely for ergonomics: it lets examples mix LayerDef and Seq in the same seq! expression without relying on Lean's coercion insertion heuristics.

• asSeq {σ τ : Spec.Shape} : F σ τ → Seq σ τ

  Convert a layer-like thing into a Seq so seq! can compose it.

@[implicit_reducible]

instance NN.API.TorchLean.NN.instAsSeqKLayerDef : AsSeqK LayerDef

A single LayerDef can always be viewed as a 1-layer sequential model (seq1).

@[implicit_reducible]

instance NN.API.TorchLean.NN.instAsSeqKSeq : AsSeqK Seq

A sequential model is already a sequential model (identity).

def NN.API.TorchLean.NN.compAny {σ τ υ : Spec.Shape} {F : Spec.Shape → Spec.Shape → Sort u} {G : Spec.Shape → Spec.Shape → Sort v} [AsSeqK F] [AsSeqK G] (f : F σ τ) (g : G τ υ) : Seq σ υ

Compose either layers or sequential models without relying on coercions.

This is the helper used by the seq! ... macro so examples can write seq! layer1, model2, layer3 while still mirroring PyTorch's "stack layers" style.

Deterministic RNG helpers re-exported for the runtime facade.

These are small utilities used by demos and training loops (keyOf, nextSeed, uniform, mask).

Sequential Layer Helpers

Runtime.Autograd.TorchLean.NN exposes layers (LayerDef σ τ) and sequential models (Seq σ τ). For demos we often want to "just stack layers", so this namespace provides small helpers that return Seq directly (and compute a few common derived shapes like flattenLinear).

For the more fully-documented public surface (named-field configs, blocks, etc.), see NN.API.Public under API.nn.

def NN.API.TorchLean.Layers.of {σ τ : Spec.Shape} (layer : NN.LayerDef σ τ) : NN.Seq σ τ

Lift a single layer into a sequential model.

def NN.API.TorchLean.Layers.linear (inDim outDim : ℕ) (seedW seedB : ℕ := 0) : NN.Seq (Tensor.Shape.Vec inDim) (Tensor.Shape.Vec outDim)

Linear layer over vectors (returns a 1-layer Seq).

def NN.API.TorchLean.Layers.relu {s : Spec.Shape} : NN.Seq s s

Elementwise ReLU.

def NN.API.TorchLean.Layers.silu {s : Spec.Shape} : NN.Seq s s

Elementwise SiLU/Swish.

def NN.API.TorchLean.Layers.gelu {s : Spec.Shape} : NN.Seq s s

Elementwise GELU.

def NN.API.TorchLean.Layers.sigmoid {s : Spec.Shape} : NN.Seq s s

Elementwise sigmoid.

def NN.API.TorchLean.Layers.tanh {s : Spec.Shape} : NN.Seq s s

Elementwise tanh.

def NN.API.TorchLean.Layers.softmax {s : Spec.Shape} : NN.Seq s s

Softmax layer.

def NN.API.TorchLean.Layers.square {s : Spec.Shape} : NN.Seq s s

Elementwise square.

def NN.API.TorchLean.Layers.sum {s : Spec.Shape} : NN.Seq s Spec.Shape.scalar

Reduce-sum to a scalar.

def NN.API.TorchLean.Layers.flatten {s : Spec.Shape} : NN.Seq s (Spec.Shape.dim s.size Spec.Shape.scalar)

Flatten any input shape into a 1D vector of length Spec.Shape.size s.

def NN.API.TorchLean.Layers.dropout {s : Spec.Shape} (p : Float) (seed : ℕ := 0) : NN.Seq s s

Dropout layer that is active in training mode and identity in eval mode.

def NN.API.TorchLean.Layers.flattenLinear {s : Spec.Shape} (outDim : ℕ) (seedW seedB : ℕ := 0) : NN.Seq s (Tensor.Shape.Vec outDim)

Flatten -> Linear head, with the input dimension computed from the input shape.

def NN.API.TorchLean.Layers.conv2d (inC outC kH kW stride padding inH inW : ℕ) {hInC : inC ≠ 0} {hKH : kH ≠ 0} {hKW : kW ≠ 0} (seedK seedB : ℕ := 0) (kInit : Runtime.Autograd.Torch.Init.Scheme := Runtime.Autograd.Torch.Init.Scheme.uniform (-0.1) 0.1) : NN.Seq (Tensor.Shape.CHW inC inH inW) (Tensor.Shape.CHW outC ((inH + 2 * padding - kH) / stride + 1) ((inW + 2 * padding - kW) / stride + 1))

Sequential 2D convolution layer for CHW inputs.

def NN.API.TorchLean.Layers.maxPool2d (kH kW inH inW inC stride : ℕ) {hKH : kH ≠ 0} {hKW : kW ≠ 0} : NN.Seq (Tensor.Shape.CHW inC inH inW) (Tensor.Shape.CHW inC ((inH - kH) / stride + 1) ((inW - kW) / stride + 1))

Sequential max-pooling layer for CHW inputs.

def NN.API.TorchLean.Layers.maxPool2dPad (kH kW inH inW inC stride padding : ℕ) {hKH : kH ≠ 0} {hKW : kW ≠ 0} : NN.Seq (Tensor.Shape.CHW inC inH inW) (Tensor.Shape.CHW inC ((inH + 2 * padding - kH) / stride + 1) ((inW + 2 * padding - kW) / stride + 1))

Sequential padded max-pooling layer for CHW inputs.

def NN.API.TorchLean.Layers.avgPool2d (kH kW inH inW inC stride : ℕ) {hKH : kH ≠ 0} {hKW : kW ≠ 0} : NN.Seq (Tensor.Shape.CHW inC inH inW) (Tensor.Shape.CHW inC ((inH - kH) / stride + 1) ((inW - kW) / stride + 1))

Sequential average-pooling layer for CHW inputs.

def NN.API.TorchLean.Layers.avgPool2dPad (kH kW inH inW inC stride padding : ℕ) {hKH : kH ≠ 0} {hKW : kW ≠ 0} : NN.Seq (Tensor.Shape.CHW inC inH inW) (Tensor.Shape.CHW inC ((inH + 2 * padding - kH) / stride + 1) ((inW + 2 * padding - kW) / stride + 1))

Sequential padded average-pooling layer for CHW inputs.

def NN.API.TorchLean.Layers.globalAvgPoolCHW (c h w : ℕ) {hC : c > 0} {hH : h > 0} {hW : w > 0} : NN.Seq (Tensor.Shape.CHW c h w) (Tensor.Shape.Vec c)

Global average-pooling over C×H×W inputs.

PyTorch analogy: torch.nn.functional.adaptive_avg_pool2d(x, output_size=1) followed by flattening the spatial axes.

def NN.API.TorchLean.Layers.globalAvgPoolNCHW (n c h w : ℕ) {hN : n > 0} {hC : c > 0} {hH : h > 0} {hW : w > 0} : NN.Seq (Tensor.Shape.NCHW n c h w) (Spec.Shape.dim n (Spec.Shape.dim c Spec.Shape.scalar))

Global average-pooling over N×C×H×W inputs.

PyTorch analogy: torch.nn.functional.adaptive_avg_pool2d(x, output_size=1) and then reshaping to (N, C).

def NN.API.TorchLean.Layers.layerNorm (batch seqLen embedDim : ℕ) {hSeq : seqLen > 0} {hEmbed : embedDim > 0} (seedGamma seedBeta : ℕ := 0) : NN.Seq (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim embedDim Spec.Shape.scalar))) (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim embedDim Spec.Shape.scalar)))

Sequence-wise layer normalization.

PyTorch analogy: torch.nn.LayerNorm(embedDim) applied to each position in a sequence.

def NN.API.TorchLean.Layers.rmsNorm (batch seqLen embedDim : ℕ) {hSeq : seqLen > 0} {hEmbed : embedDim > 0} (seedGamma : ℕ := 0) : NN.Seq (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim embedDim Spec.Shape.scalar))) (Spec.Shape.dim batch (Spec.Shape.dim seqLen (Spec.Shape.dim embedDim Spec.Shape.scalar)))

Sequence-wise RMS normalization.

PyTorch analogy: an RMSNorm-style layer over (seqLen × embedDim) tensors.

def NN.API.TorchLean.Layers.batchNormCHW (channels height width : ℕ) {hC : channels > 0} {hH : height > 0} {hW : width > 0} (seedGamma seedBeta seedMean seedVar : ℕ := 0) : NN.Seq (Tensor.Shape.CHW channels height width) (Tensor.Shape.CHW channels height width)

Mode-aware batch norm on a single C×H×W image tensor.

PyTorch analogy: torch.nn.BatchNorm2d(channels) on a single sample, with the layer's mode controlling whether running statistics are updated or reused.

def NN.API.TorchLean.Layers.batchNormEvalCHW (channels height width : ℕ) {hC : channels > 0} {hH : height > 0} {hW : width > 0} (seedGamma seedBeta seedMean seedVar : ℕ := 0) : NN.Seq (Tensor.Shape.CHW channels height width) (Tensor.Shape.CHW channels height width)

Eval-mode batch norm on a single C×H×W image tensor with explicit running statistics.

PyTorch analogy: torch.nn.BatchNorm2d(...).eval() with running_mean and running_var.

def NN.API.TorchLean.Layers.instanceNorm2dNCHW (n c h w : ℕ) {hN : n > 0} {hC : c > 0} {hH : h > 0} {hW : w > 0} (seedGamma seedBeta : ℕ := 0) : NN.Seq (Tensor.Shape.NCHW n c h w) (Tensor.Shape.NCHW n c h w)

Instance normalization over N×C×H×W tensors.

PyTorch analogy: torch.nn.InstanceNorm2d(c, affine=True) with NCHW layout.

def NN.API.TorchLean.Layers.groupNorm2dNCHW (n c h w groups : ℕ) {hN : n > 0} {hC : c > 0} {hH : h > 0} {hW : w > 0} {hG : groups > 0} (hGE : c ≥ groups) (hDiv : c % groups = 0) (seedGamma seedBeta : ℕ := 0) : NN.Seq (Tensor.Shape.NCHW n c h w) (Tensor.Shape.NCHW n c h w)

Group normalization over N×C×H×W tensors.

PyTorch analogy: torch.nn.GroupNorm(groups, c) with NCHW layout.

def NN.API.TorchLean.Layers.batchNorm2dNCHW (n c h w : ℕ) {hN : n > 0} {hC : c > 0} {hH : h > 0} {hW : w > 0} (seedGamma seedBeta seedMean seedVar : ℕ := 0) : NN.Seq (Tensor.Shape.NCHW n c h w) (Tensor.Shape.NCHW n c h w)

Batch norm over N×C×H×W tensors in training mode.

PyTorch analogy: torch.nn.BatchNorm2d(c) during training, where batch statistics are used.

def NN.API.TorchLean.Layers.attention (batch n dModel numHeads headDim : ℕ) {hN : n ≠ 0} (seedW : ℕ := 0) (mask : Option (Spec.Tensor Bool (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) := none) : NN.Seq (Spec.Shape.dim batch (Spec.Shape.dim n (Spec.Shape.dim dModel Spec.Shape.scalar))) (Spec.Shape.dim batch (Spec.Shape.dim n (Spec.Shape.dim dModel Spec.Shape.scalar)))

Multi-head self-attention over sequence embeddings.

PyTorch analogy: torch.nn.MultiheadAttention(embed_dim=dModel, num_heads=numHeads) in self- attention mode, with explicit n × dModel shapes.

@[reducible, inline]

abbrev NN.API.TorchLean.Autodiff.Model.Params {σ τ : Spec.Shape} (model : NN.Seq σ τ) (α : Type) : Type

Parameter list type for a given model (a TList over Seq.paramShapes).

@[reducible, inline]

abbrev NN.API.TorchLean.Autodiff.Model.OutputLoss (τ υ : Spec.Shape) : Type 1

Loss function over a model output and a target.

This is expressed in terms of RefTy so it works uniformly for eager execution and compiled execution.

def NN.API.TorchLean.Autodiff.Model.initParamsWith {σ τ : Spec.Shape} (model : NN.Seq σ τ) {α : Type} (cast : Float → α) : Params model α

Initialize model parameters by casting the model's Float initializers elementwise using cast.

def NN.API.TorchLean.Autodiff.Model.initParams {σ τ : Spec.Shape} (model : NN.Seq σ τ) {α : Type} [Runtime.Scalar α] : Params model α

Initialize model parameters using the runtime literal injection API.Runtime.ofFloat.

def NN.API.TorchLean.Autodiff.Model.linearParams {α : Type} {inDim outDim seedW seedB : ℕ} (w : Spec.Tensor α (Tensor.Shape.Mat outDim inDim)) (b : Spec.Tensor α (Tensor.Shape.Vec outDim)) : Params (Layers.linear inDim outDim seedW seedB) α

Pack explicit weight and bias tensors for a single Layers.linear model.

def NN.API.TorchLean.Autodiff.Model.OutputLoss.mse {τ : Spec.Shape} (reduction : Loss.Reduction := Runtime.Autograd.TorchLean.Loss.Reduction.mean) : OutputLoss τ τ

Mean-squared error loss (mse) between yhat and y.

def NN.API.TorchLean.Autodiff.Model.OutputLoss.crossEntropyOneHot {τ : Spec.Shape} (reduction : Loss.Reduction := Runtime.Autograd.TorchLean.Loss.Reduction.mean) : OutputLoss τ τ

Cross-entropy loss between logits and one-hot targets. PyTorch analogue: nn.CrossEntropyLoss.

def NN.API.TorchLean.Autodiff.Model.OutputLoss.detach {τ υ : Spec.Shape} (loss : OutputLoss τ υ) : OutputLoss τ υ

Detach the model output before feeding it into a loss.

This is useful when you want to compute a metric loss without backpropagating through it.

def NN.API.TorchLean.Autodiff.Model.lossProgram {σ τ υ : Spec.Shape} (model : NN.Seq σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] : Program α (model.paramShapes ++ [σ, υ]) Spec.Shape.scalar

Build a TorchLean Program that computes a scalar loss from (params, x, target).

This is the bridge between Seq.program (which produces model outputs) and the autograd entry points (which expect a scalar-valued program).

def NN.API.TorchLean.Autodiff.Model.vjpParams {σ τ : Spec.Shape} (model : NN.Seq σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (seedOut : Spec.Tensor α τ) : IO (Params model α)

VJP of the model output w.r.t. parameters.

def NN.API.TorchLean.Autodiff.Model.vjpInputs {σ τ : Spec.Shape} (model : NN.Seq σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (seedOut : Spec.Tensor α τ) : IO (TList α [σ])

VJP of the model output w.r.t. inputs.

def NN.API.TorchLean.Autodiff.Model.jacrevParams {σ τ : Spec.Shape} (model : NN.Seq σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) : IO (Array (Params model α))

Jacobian (reverse-mode) of the model output w.r.t. parameters, returned as rows.

def NN.API.TorchLean.Autodiff.Model.gradParams {σ τ υ : Spec.Shape} (model : NN.Seq σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (Params model α)

Gradient of loss(model(params, x), target) w.r.t. parameters.

def NN.API.TorchLean.Autodiff.Model.gradInputs {σ τ υ : Spec.Shape} (model : NN.Seq σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) : IO (TList α [σ, υ])

Gradient of loss(model(params, x), target) w.r.t. inputs (x and target).

def NN.API.TorchLean.Autodiff.Model.jvpParams {σ τ υ : Spec.Shape} (model : NN.Seq σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) (vparams : Params model α) : IO α

JVP of a scalar loss w.r.t. parameters in direction vparams.

def NN.API.TorchLean.Autodiff.Model.hvpParams {σ τ υ : Spec.Shape} (model : NN.Seq σ τ) (loss : OutputLoss τ υ) {α : Type} [Context α] [DecidableEq Spec.Shape] (params : Params model α) (x : Spec.Tensor α σ) (target : Spec.Tensor α υ) (vparams : Params model α) : IO (Params model α)

HVP (Hessian-vector product) of a scalar loss w.r.t. parameters in direction vparams.

@[reducible, inline]

abbrev NN.API.TorchLean.Autodiff.Function1.Fn (σ τ : Spec.Shape) : Type 1

Type of a pure tensor function expressed in RefTy form.

This matches the calling convention expected by TorchLean.Program/autodiff compilation.

def NN.API.TorchLean.Autodiff.Function1.program {σ τ : Spec.Shape} (f : Fn σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] : Program α [σ] τ

Turn an Fn into a single-input TorchLean Program.

def NN.API.TorchLean.Autodiff.Function1.jacfwd {σ τ : Spec.Shape} (f : Fn σ τ) {α : Type} [Context α] [DecidableEq Spec.Shape] (x : Spec.Tensor α σ) : IO (Array (Spec.Tensor α τ))

Forward-mode Jacobian (rows) of a pure function.

def NN.API.TorchLean.Autodiff.Function1.hessian {σ : Spec.Shape} (f : Fn σ Spec.Shape.scalar) {α : Type} [Context α] [DecidableEq Spec.Shape] (x : Spec.Tensor α σ) : IO (Array (Spec.Tensor α σ))

Hessian for a scalar-valued pure function.

Model constructors (re-export)

This namespace re-exports a small set of ready-made model constructors (MLP/CNN/ResNet18/etc.), primarily for runnable demos and smoke tests.

For compositional building blocks, prefer API.TorchLean.NN and API.TorchLean.Layers.

ScalarModule API (Session-Like Interface)

The ScalarModule interface is the TorchLean equivalent of "instantiate a model, then do forward, backward, and optimizer steps" in an imperative runtime.

This section mostly re-exports Runtime.Autograd.TorchLean.Module.* and adds small CLI-friendly helpers (Module.withModule / Module.withModuleRuntime) that select dtype/backend from flags.

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

Instantiate a ScalarModuleDef under explicit Torch options (backend, fastKernels, useGpu, etc.).

This is the most direct "runtime" entrypoint (used by the CPU/CUDA example binaries), since it threads the same options record all the way down to the eager tape / CUDA tape selection.

Note: instantiate (without options) is a convenience wrapper that only selects the backend and uses default runtime options.

structure NN.API.TorchLean.Module.ExecConfig : Type

Execution configuration parsed from CLI flags.

Supported flags (parsed by ExecConfig.parseAndStrip):

• --dtype ... / --float32-mode ... (see NN.API.DType)
• --backend eager|compiled
• --cpu / --cuda (eager device selection)
• --fast-kernels (eager-only performance hooks, no effect on compiled backend)
• --fast-gpu-matmul-precision fp32|fp64 (fast-kernel CUDA matmul precision)

• dtype : DType

  Scalar dtype selection.

• backend : Backend

  Execution backend selection.

• useGpu : Bool

  Eager execution device selector.

  When true and backend = .eager, TorchLean uses the CUDA tape backend.

• fastKernels : Bool

  Enable runtime-only eager fast kernels (tight-loop implementations for a few hot ops).

• fastGpuMatmulPrecision : GpuMatmulPrecision

  GPU precision for fast-kernel matmul over Lean Float tensors.

@[implicit_reducible]

instance NN.API.TorchLean.Module.instReprExecConfig : Repr ExecConfig

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

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

@[implicit_reducible]

instance NN.API.TorchLean.Module.instDecidableEqExecConfig : DecidableEq ExecConfig

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

Parse a backend selector string into a runtime Backend.

def NN.API.TorchLean.Module.ExecConfig.parseFastGpuMatmulPrecision (v : String) : Except String GpuMatmulPrecision

Parse a fast-kernel CUDA matmul precision selector.

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

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

Consumed flags:

• --backend eager|compiled (at most once),
• --cpu / --cuda (boolean flags; last one wins; removed from rest),
• --fast-kernels (boolean flag; removed from rest).
• --fast-gpu-matmul-precision fp32|fp64 (at most once).

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

Default dtype policy:

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

def NN.API.TorchLean.Module.ExecConfig.parseAndStripWithDefaultDType.go (useGpu fastKernels : Bool) (acc : List String) : List String → Bool × Bool × List String

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

Parse CLI flags with the standard TorchLean default dtype policy.

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

Log the chosen execution config to stdout (for reproducible demos).

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

Parse runtime flags (--dtype, --backend, --cpu|--cuda, --fast-kernels, --fast-gpu-matmul-precision) 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 NN.API.TorchLean.Module.withModule {paramShapes inputShapes : List Spec.Shape} (defn : ScalarModuleDef paramShapes inputShapes) (args : List String) (k : {α : Type} → [inst : Semantics.Scalar α] → [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, --cpu|--cuda, --fast-kernels, --fast-gpu-matmul-precision), then call a continuation.

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

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

Like withModule, but also provides an API.Runtime.Scalar α instance (for numeric literals).

Executable main Helpers

TorchLean has a lot of pure, type-indexed code (models live in Type 2), but runnable scripts still want a "single entrypoint" that handles:

• parsing --seed,
• selecting an executable dtype/backend/device from flags,
• seeding TorchLean's global RNG stream (API.rand) so nn.freshSeed/nn.withModel are deterministic.

structure NN.API.TorchLean.Module.RunOptions : Type

Options for TorchLean.Module.run (banner printing, trailing ok, etc.).

• banner? : Option (Options → String)

  Optional banner to print before executing the program.

• flush : Bool

  Flush stdout after printing the banner (if present).

• printOk : Bool

  Print "{exeName}: ok" on success.

@[implicit_reducible]

instance NN.API.TorchLean.Module.instInhabitedRunOptions : Inhabited RunOptions

def NN.API.TorchLean.Module.instInhabitedRunOptions.default : RunOptions

def NN.API.TorchLean.Module.RunOptions.printBanner (o : RunOptions) (opts : Options) : IO Unit

inductive NN.API.TorchLean.Module.RunAction : Type 1

How run should select the scalar backend for an executable.

• any (k : {α : Type} → [Semantics.Scalar α] → [DecidableEq Spec.Shape] → [ToString α] → [Runtime.Scalar α] → (Float → α) → Options → List String → IO Unit) : RunAction

  Allow --dtype selection (the continuation must work for all executable scalar backends).

• float (k : Options → List String → IO Unit) : RunAction

  Force the scalar backend to builtin Float (useful for Float-only IO bridges / CUDA upload paths).

def NN.API.TorchLean.Module.run (exeName : String) (args : List String) (action : RunAction) (runOpts : RunOptions := { }) : IO UInt32

CLI entrypoint helper for executable main functions.

This parses:

• --seed N (via API.CLI.takeSeed), and
• runtime execution flags (--dtype, --float32-mode, --backend, --cpu|--cuda, --fast-kernels, --fast-gpu-matmul-precision), then executes the chosen RunAction.

It also seeds TorchLean's global RNG stream (API.rand) so code that draws init seeds via API.nn.freshSeed/API.nn.withModel is deterministic by default, matching the PyTorch pattern of calling torch.manual_seed once in main.

inductive NN.API.TorchLean.Supervised.SeqLoss : Type

Built-in loss choices for SeqTask.

• mse (reduction : Loss.Reduction := Runtime.Autograd.TorchLean.Loss.Reduction.mean) : SeqLoss
• crossEntropyOneHot (reduction : Loss.Reduction := Runtime.Autograd.TorchLean.Loss.Reduction.mean) : SeqLoss

structure NN.API.TorchLean.Supervised.SeqTask (σ τ : Spec.Shape) : Type 2

A supervised task is just a model plus a choice of loss.

• model : NN.Seq σ τ

  Model to run.

• loss : SeqLoss

  Loss function.

def NN.API.TorchLean.Supervised.SeqTask.moduleDefWithMode {σ τ : Spec.Shape} (task : SeqTask σ τ) (mode : NN.Mode) : Runtime.Autograd.TorchLean.ScalarModuleDef task.model.paramShapes [σ, τ]

Build a ScalarModuleDef for a task, choosing an explicit model mode (train/eval).

This is the underlying "instantiate me as a runnable module" step for training.

def NN.API.TorchLean.Supervised.SeqTask.moduleDef {σ τ : Spec.Shape} (task : SeqTask σ τ) : Runtime.Autograd.TorchLean.ScalarModuleDef task.model.paramShapes [σ, τ]

Default module definition for a task (training mode).

def NN.API.TorchLean.Supervised.SeqTask.mse {σ τ : Spec.Shape} (model : NN.Seq σ τ) (reduction : Loss.Reduction := Runtime.Autograd.TorchLean.Loss.Reduction.mean) : SeqTask σ τ

Constructor: regression task (MSE loss).

def NN.API.TorchLean.Supervised.SeqTask.crossEntropyOneHot {σ τ : Spec.Shape} (model : NN.Seq σ τ) (reduction : Loss.Reduction := Runtime.Autograd.TorchLean.Loss.Reduction.mean) : SeqTask σ τ

Constructor: one-hot classification task (cross-entropy loss).

@[reducible, inline]

abbrev NN.API.TorchLean.Supervised.paramShapes {σ τ : Spec.Shape} (task : SeqTask σ τ) : List Spec.Shape

Parameter shapes for a task (delegates to Seq.paramShapes).

inductive NN.API.TorchLean.Supervised.OptimizerConfig : Type

Optimizer hyperparameter configuration for the supervised training helpers.

We keep this small for examples and lightweight trainers. It mirrors a few common PyTorch optimizers by name/defaults, but it does not try to cover the full option surface of torch.optim.*.

• sgd (lr : Float) (momentum : Float := 0.0) : OptimizerConfig

  SGD optimizer config.

  PyTorch analogy: torch.optim.SGD(..., lr=..., momentum=...) when momentum > 0, and plain SGD when momentum = 0.

• adam (lr : Float) (beta1 : Float := 0.9) (beta2 : Float := 0.999) (epsilon : Float := 1e-8) : OptimizerConfig

  Adam optimizer config.

• adamw (lr : Float) (weightDecay : Float := 1e-2) (beta1 : Float := 0.9) (beta2 : Float := 0.999) (epsilon : Float := 1e-8) : OptimizerConfig

  AdamW optimizer config (decoupled weight decay).

@[implicit_reducible]

instance NN.API.TorchLean.Supervised.instReprOptimizerConfig : Repr OptimizerConfig

def NN.API.TorchLean.Supervised.instReprOptimizerConfig.repr : OptimizerConfig → ℕ → Std.Format

structure NN.API.TorchLean.Supervised.FitConfig : Type

Step-based training configuration for fit / fitDataset.

Fields:

• steps: number of parameter updates,
• optimizer: optimizer hyperparameters,
• scheduler: optional learning-rate schedule (applied per step),
• logEvery: progress printing frequency (0 disables logging).

• steps : ℕ

  Number of training steps.

• optimizer : OptimizerConfig

  Optimizer configuration.

• scheduler : Option Schedulers.Config

  Scheduler configuration.

• logEvery : ℕ

  Log once every this many steps.

def NN.API.TorchLean.Supervised.instReprFitConfig.repr : FitConfig → ℕ → Std.Format

@[implicit_reducible]

instance NN.API.TorchLean.Supervised.instReprFitConfig : Repr FitConfig

structure NN.API.TorchLean.Supervised.FitReport (α : Type) : Type

Small summary returned by fit* helpers.

By default, before and after are mean loss values, but the type is polymorphic so callers can report other scalars in the same shape.

• before : α

  Metrics before training.

• after : α

  Metrics after training.

structure NN.API.TorchLean.Supervised.LoaderFitConfig : Type

Epoch-based training configuration for fitLoader (data-loader training).

Fields:

• epochs: number of epochs (each epoch iterates once over the loader),
• optimizer: optimizer hyperparameters,
• scheduler: optional learning-rate schedule (applied per step/epoch depending on helper),
• logEvery: progress printing frequency (0 disables logging).

• epochs : ℕ

  Number of epochs to train for.

• optimizer : OptimizerConfig

  Optimizer configuration.

• scheduler : Option Schedulers.Config

  Scheduler configuration.

• logEvery : ℕ

  Log once every this many steps.

@[implicit_reducible]

instance NN.API.TorchLean.Supervised.instReprLoaderFitConfig : Repr LoaderFitConfig

def NN.API.TorchLean.Supervised.instReprLoaderFitConfig.repr : LoaderFitConfig → ℕ → Std.Format

def NN.API.TorchLean.Supervised.optimizerLR : OptimizerConfig → Float

Extract the base learning rate encoded in an optimizer configuration.

def NN.API.TorchLean.Supervised.stepLR (scheduler : Option Schedulers.Config) (cfg : OptimizerConfig) (step : ℕ) : Float

Resolve the learning rate to use at a given training step.

If a scheduler is present, it takes precedence over the optimizer's baked-in base learning rate. Otherwise this simply returns optimizerLR cfg.

def NN.API.TorchLean.Supervised.mapStateList {State : Type → Spec.Shape → Type} {α : Type} {ss : List Spec.Shape} : ({s : Spec.Shape} → State α s → State α s) → Runtime.Autograd.TorchLean.StateList State α ss → Runtime.Autograd.TorchLean.StateList State α ss

Map a state update over every optimizer-state entry in a shape-indexed parameter list.

def NN.API.TorchLean.Supervised.adamStateWithLR {α : Type} (lr : α) {paramShapes : List Spec.Shape} : Runtime.Autograd.TorchLean.StateList Optim.Adam.State α paramShapes → Runtime.Autograd.TorchLean.StateList Optim.Adam.State α paramShapes

Set the learning rate field of every Adam optimizer state entry to lr.

def NN.API.TorchLean.Supervised.momentumSGDStateWithLR {α : Type} (lr : α) {paramShapes : List Spec.Shape} : Runtime.Autograd.TorchLean.StateList Optim.MomentumSGD.State α paramShapes → Runtime.Autograd.TorchLean.StateList Optim.MomentumSGD.State α paramShapes

Set the learning rate field of every momentum-SGD optimizer state entry to lr.

def NN.API.TorchLean.Supervised.adamwStateWithLR {α : Type} (lr : α) {paramShapes : List Spec.Shape} : Runtime.Autograd.TorchLean.StateList Optim.AdamW.State α paramShapes → Runtime.Autograd.TorchLean.StateList Optim.AdamW.State α paramShapes

Set the learning rate field of every AdamW optimizer state entry to lr.

structure NN.API.TorchLean.Supervised.Runner (α : Type) [Semantics.Scalar α] [DecidableEq Spec.Shape] {σ τ : Spec.Shape} (task : SeqTask σ τ) : Type

A fully instantiated supervised task runner.

This bundles:

• the imperative ScalarModule (parameters/buffers stored in refs),
• compiled predictors and loss functions for both .train and .eval modes (so switching mode is cheap),
• and the current mode stored in an IO.Ref.

The mode influences both operator behavior (e.g. dropout/batchnorm) and whether buffers are updated during training.

• module : Runtime.Autograd.TorchLean.ScalarModule α (paramShapes task) [σ, τ]

  Instantiated scalar module storing parameters/buffers in mutable refs.

• predictorTrain : CompiledOut α (paramShapes task ++ [σ]) τ

  Compiled forward predictor specialized to training-mode behavior.

• predictorEval : CompiledOut α (paramShapes task ++ [σ]) τ

  Compiled forward predictor specialized to eval-mode behavior.

• lossTrain : CompiledScalar α (paramShapes task ++ [σ, τ])

  Compiled loss function for training-mode behavior.

• lossEval : CompiledScalar α (paramShapes task ++ [σ, τ])

  Compiled loss function for eval-mode behavior.

• mode : IO.Ref NN.Mode

  Mutable mode flag (.train / .eval) used by stateful layers (e.g. dropout/batchnorm).

def NN.API.TorchLean.Supervised.instantiateWithOptions {σ τ : Spec.Shape} (task : SeqTask σ τ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [Runtime.Autograd.Torch.Internal.CudaBridge.TensorConv α] (cast : Float → α) (opts : Options := { }) : IO (Runner α task)

Instantiate a Runner by explicitly providing a Float → α cast and backend.

Use this when you want to run the same task over different numeric backends (e.g. Float vs IEEE32Exec) or when you want custom literal injection.

def NN.API.TorchLean.Supervised.instantiateWith {σ τ : Spec.Shape} (task : SeqTask σ τ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [Runtime.Autograd.Torch.Internal.CudaBridge.TensorConv α] (cast : Float → α) (backend : Backend := Runtime.Autograd.Torch.Backend.eager) : IO (Runner α task)

Instantiate a Runner by explicitly providing a Float → α cast and a backend selector.

This is a backward-compatible wrapper over instantiateWithOptions.

def NN.API.TorchLean.Supervised.instantiateWithRuntimeOptions {σ τ : Spec.Shape} (task : SeqTask σ τ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [Runtime.Scalar α] [Runtime.Autograd.Torch.Internal.CudaBridge.TensorConv α] (opts : Options := { }) : IO (Runner α task)

Instantiate a Runner using the standard runtime literal injection API.Runtime.ofFloat.

This is the common entrypoint for executable examples.

def NN.API.TorchLean.Supervised.instantiate {σ τ : Spec.Shape} (task : SeqTask σ τ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [Runtime.Scalar α] [Runtime.Autograd.Torch.Internal.CudaBridge.TensorConv α] (backend : Backend := Runtime.Autograd.Torch.Backend.eager) : IO (Runner α task)

Instantiate a Runner using the standard runtime literal injection API.Runtime.ofFloat and a backend selector.

This is a backward-compatible wrapper over instantiateWithRuntimeOptions.

def NN.API.TorchLean.Supervised.run {σ τ : Spec.Shape} (task : SeqTask σ τ) (args : List String) (k : {α : Type} → [inst : Semantics.Scalar α] → [inst_1 : DecidableEq Spec.Shape] → [ToString α] → [Runtime.Scalar α] → Runner α task → List String → IO Unit) : IO Unit

Run a TorchLean task with CLI-style dtype/backend selection, then call k with a fully constructed runner.

This is used by lake exe entrypoints: run takes care of parsing dtype flags and instantiating the underlying module/compiled programs.

def NN.API.TorchLean.Supervised.params {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO (TList α (paramShapes task))

Read the current parameter list from a runner.

def NN.API.TorchLean.Supervised.mode {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO NN.Mode

Read the runner's current mode (.train or .eval).

def NN.API.TorchLean.Supervised.setMode {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) (value : NN.Mode) : IO Unit

Set the runner mode (.train or .eval).

def NN.API.TorchLean.Supervised.trainMode {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO Unit

Convenience: setMode runner .train.

def NN.API.TorchLean.Supervised.evalMode {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO Unit

Convenience: setMode runner .eval.

def NN.API.TorchLean.Supervised.isTraining {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO Bool

Predicate: are we in training mode?

def NN.API.TorchLean.Supervised.activePredictor {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO (CompiledOut α (paramShapes task ++ [σ]) τ)

Pick the predictor compiled for the runner's current mode (.train or .eval).

def NN.API.TorchLean.Supervised.activeLoss {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO (CompiledScalar α (paramShapes task ++ [σ, τ]))

Pick the loss program compiled for the runner's current mode (.train or .eval).

def NN.API.TorchLean.Supervised.updateRunnerBuffers {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) (sample : TList α [σ, τ]) : IO Unit

Refresh mode-dependent runner buffers using one supervised sample.

This mutates the module parameters only in .train mode, mirroring PyTorch-style buffer updates for layers such as normalization. In .eval mode it is a no-op.

def NN.API.TorchLean.Supervised.backward {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) (sample : TList α [σ, τ]) : IO (TList α (paramShapes task))

Run one forward/backward pass on a single supervised sample and return gradients for all parameters.

This is the TorchLean analogue of the loss.backward() payload in PyTorch, except TorchLean returns the gradients explicitly instead of storing them in .grad fields.

def NN.API.TorchLean.Supervised.predict {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) (x : Spec.Tensor α σ) : IO (Spec.Tensor α τ)

Predict on one input tensor using the runner's active mode (.train or .eval).

def NN.API.TorchLean.Supervised.predictBatch {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) (xs : List (Spec.Tensor α σ)) : IO (List (Spec.Tensor α τ))

Predict on a list of inputs by repeatedly calling predict.

def NN.API.TorchLean.Supervised.predictClass? {σ : Spec.Shape} {n : ℕ} {task : SeqTask σ (Spec.Shape.dim n Spec.Shape.scalar)} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (runner : Runner α task) (x : Spec.Tensor α σ) : IO (Option (Fin n))

For classification heads: run predict, then take argmax over the logits (if defined).

def NN.API.TorchLean.Supervised.accuracyOneHot {σ : Spec.Shape} {n : ℕ} {task : SeqTask σ (Spec.Shape.dim n Spec.Shape.scalar)} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (runner : Runner α task) (samples : List (TList α [σ, Spec.Shape.dim n Spec.Shape.scalar])) : IO (ℕ × ℕ)

Compute (correct, total) for a one-hot classification dataset.

def NN.API.TorchLean.Supervised.accuracyOneHot.go {σ : Spec.Shape} {n : ℕ} {task : SeqTask σ (Spec.Shape.dim n Spec.Shape.scalar)} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (runner : Runner α task) (correct total : ℕ) : List (TList α [σ, Spec.Shape.dim n Spec.Shape.scalar]) → IO (ℕ × ℕ)

def NN.API.TorchLean.Supervised.meanLoss {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] (runner : Runner α task) (samples : List (TList α [σ, τ])) : IO α

Mean scalar loss over a list of supervised samples (uses the runner's active mode).

def NN.API.TorchLean.Supervised.meanLossDataset {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] (runner : Runner α task) (dataset : Runtime.Autograd.Train.Dataset (TList α [σ, τ])) : IO α

Mean scalar loss over a dataset (materialized via dataset.toList).

def NN.API.TorchLean.Supervised.moduleLoss {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) (sample : TList α [σ, τ]) : IO α

Scalar loss for one sample through the instantiated runtime module.

def NN.API.TorchLean.Supervised.fit {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] [Runtime.Scalar α] (runner : Runner α task) (cfg : FitConfig) (samples : List (TList α [σ, τ])) : IO (FitReport α)

Fit on a small in-memory list of supervised samples for a fixed number of steps.

This is the simplest training-loop helper: it is intended for examples and small synthetic datasets. For loader-based training, see fitLoader.

def NN.API.TorchLean.Supervised.fitDataset {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] [Runtime.Scalar α] (runner : Runner α task) (cfg : FitConfig) (dataset : Runtime.Autograd.Train.Dataset (TList α [σ, τ])) : IO (FitReport α)

fit over a dataset (materialized as a list).

def NN.API.TorchLean.Supervised.fitLoader {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] [Runtime.Scalar α] (runner : Runner α task) (cfg : LoaderFitConfig) (dl : Runtime.Autograd.Train.DataLoader (TList α [σ, τ])) : IO (FitReport α × Runtime.Autograd.Train.DataLoader (TList α [σ, τ]))

Fit over a DataLoader for cfg.epochs epochs, returning the final report and the updated loader.

This corresponds to the common PyTorch pattern: for epoch in ...: for batch in loader: step(batch).

def NN.API.TorchLean.Supervised.fitLoader.trainSgdBatches {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [Runtime.Scalar α] (runner : Runner α task) (cfg : LoaderFitConfig) (epoch : ℕ) (batches : List (List (TList α [σ, τ]))) : IO Unit

def NN.API.TorchLean.Supervised.fitLoader.runSgdEpochs {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [Runtime.Scalar α] (runner : Runner α task) (cfg : LoaderFitConfig) (nextEpoch : Runtime.Autograd.Train.DataLoader (TList α [σ, τ]) → IO (Runtime.Autograd.Train.DataLoader (TList α [σ, τ]) × List (List (TList α [σ, τ])))) (remaining epoch : ℕ) (loader : Runtime.Autograd.Train.DataLoader (TList α [σ, τ])) : IO (Runtime.Autograd.Train.DataLoader (TList α [σ, τ]))

def NN.API.TorchLean.Supervised.fitLoader.trainMomSamples {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Runtime.Scalar α] (runner : Runner α task) (cfg : LoaderFitConfig) (lr momentum : Float) (epoch stepIdx : ℕ) (state : (Runtime.Autograd.TorchLean.momentumSGD (Runtime.ofFloat lr) (Runtime.ofFloat momentum)).State) (samples : List (TList α [σ, τ])) : have opt := Runtime.Autograd.TorchLean.momentumSGD (Runtime.ofFloat lr) (Runtime.ofFloat momentum); IO (opt.State × ℕ)

def NN.API.TorchLean.Supervised.fitLoader.trainMomBatches {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Runtime.Scalar α] (runner : Runner α task) (cfg : LoaderFitConfig) (lr momentum : Float) (epoch stepIdx : ℕ) (state : (Runtime.Autograd.TorchLean.momentumSGD (Runtime.ofFloat lr) (Runtime.ofFloat momentum)).State) (batches : List (List (TList α [σ, τ]))) : have opt := Runtime.Autograd.TorchLean.momentumSGD (Runtime.ofFloat lr) (Runtime.ofFloat momentum); IO (opt.State × ℕ)

def NN.API.TorchLean.Supervised.fitLoader.runMomEpochs {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Runtime.Scalar α] (runner : Runner α task) (cfg : LoaderFitConfig) (nextEpoch : Runtime.Autograd.Train.DataLoader (TList α [σ, τ]) → IO (Runtime.Autograd.Train.DataLoader (TList α [σ, τ]) × List (List (TList α [σ, τ])))) (lr momentum : Float) (epoch remaining : ℕ) (loader : Runtime.Autograd.Train.DataLoader (TList α [σ, τ])) (st : (Runtime.Autograd.TorchLean.momentumSGD (Runtime.ofFloat lr) (Runtime.ofFloat momentum)).State) : have opt := Runtime.Autograd.TorchLean.momentumSGD (Runtime.ofFloat lr) (Runtime.ofFloat momentum); IO (Runtime.Autograd.Train.DataLoader (TList α [σ, τ]) × opt.State)

def NN.API.TorchLean.Supervised.fitLoader.trainAdamSamples {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Runtime.Scalar α] (runner : Runner α task) (cfg : LoaderFitConfig) (lr beta1 beta2 epsilon : Float) (epoch stepIdx : ℕ) (state : (Runtime.Autograd.TorchLean.adam (Runtime.ofFloat lr) (Runtime.ofFloat beta1) (Runtime.ofFloat beta2) (Runtime.ofFloat epsilon)).State) (samples : List (TList α [σ, τ])) : have opt := Runtime.Autograd.TorchLean.adam (Runtime.ofFloat lr) (Runtime.ofFloat beta1) (Runtime.ofFloat beta2) (Runtime.ofFloat epsilon); IO (opt.State × ℕ)

def NN.API.TorchLean.Supervised.fitLoader.trainAdamBatches {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Runtime.Scalar α] (runner : Runner α task) (cfg : LoaderFitConfig) (lr beta1 beta2 epsilon : Float) (epoch stepIdx : ℕ) (state : (Runtime.Autograd.TorchLean.adam (Runtime.ofFloat lr) (Runtime.ofFloat beta1) (Runtime.ofFloat beta2) (Runtime.ofFloat epsilon)).State) (batches : List (List (TList α [σ, τ]))) : have opt := Runtime.Autograd.TorchLean.adam (Runtime.ofFloat lr) (Runtime.ofFloat beta1) (Runtime.ofFloat beta2) (Runtime.ofFloat epsilon); IO (opt.State × ℕ)

def NN.API.TorchLean.Supervised.fitLoader.runAdamEpochs {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Runtime.Scalar α] (runner : Runner α task) (cfg : LoaderFitConfig) (nextEpoch : Runtime.Autograd.Train.DataLoader (TList α [σ, τ]) → IO (Runtime.Autograd.Train.DataLoader (TList α [σ, τ]) × List (List (TList α [σ, τ])))) (lr beta1 beta2 epsilon : Float) (epoch remaining : ℕ) (loader : Runtime.Autograd.Train.DataLoader (TList α [σ, τ])) (st : (Runtime.Autograd.TorchLean.adam (Runtime.ofFloat lr) (Runtime.ofFloat beta1) (Runtime.ofFloat beta2) (Runtime.ofFloat epsilon)).State) : have opt := Runtime.Autograd.TorchLean.adam (Runtime.ofFloat lr) (Runtime.ofFloat beta1) (Runtime.ofFloat beta2) (Runtime.ofFloat epsilon); IO (Runtime.Autograd.Train.DataLoader (TList α [σ, τ]) × opt.State)

def NN.API.TorchLean.Supervised.fitLoader.trainAdamWSamples {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Runtime.Scalar α] (runner : Runner α task) (cfg : LoaderFitConfig) (lr weightDecay beta1 beta2 epsilon : Float) (epoch stepIdx : ℕ) (state : (Runtime.Autograd.TorchLean.adamw (Runtime.ofFloat lr) (Runtime.ofFloat weightDecay) (Runtime.ofFloat beta1) (Runtime.ofFloat beta2) (Runtime.ofFloat epsilon)).State) (samples : List (TList α [σ, τ])) : have opt := Runtime.Autograd.TorchLean.adamw (Runtime.ofFloat lr) (Runtime.ofFloat weightDecay) (Runtime.ofFloat beta1) (Runtime.ofFloat beta2) (Runtime.ofFloat epsilon); IO (opt.State × ℕ)

def NN.API.TorchLean.Supervised.fitLoader.trainAdamWBatches {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Runtime.Scalar α] (runner : Runner α task) (cfg : LoaderFitConfig) (lr weightDecay beta1 beta2 epsilon : Float) (epoch stepIdx : ℕ) (state : (Runtime.Autograd.TorchLean.adamw (Runtime.ofFloat lr) (Runtime.ofFloat weightDecay) (Runtime.ofFloat beta1) (Runtime.ofFloat beta2) (Runtime.ofFloat epsilon)).State) (batches : List (List (TList α [σ, τ]))) : have opt := Runtime.Autograd.TorchLean.adamw (Runtime.ofFloat lr) (Runtime.ofFloat weightDecay) (Runtime.ofFloat beta1) (Runtime.ofFloat beta2) (Runtime.ofFloat epsilon); IO (opt.State × ℕ)

def NN.API.TorchLean.Supervised.fitLoader.runAdamWEpochs {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Runtime.Scalar α] (runner : Runner α task) (cfg : LoaderFitConfig) (nextEpoch : Runtime.Autograd.Train.DataLoader (TList α [σ, τ]) → IO (Runtime.Autograd.Train.DataLoader (TList α [σ, τ]) × List (List (TList α [σ, τ])))) (lr weightDecay beta1 beta2 epsilon : Float) (epoch remaining : ℕ) (loader : Runtime.Autograd.Train.DataLoader (TList α [σ, τ])) (st : (Runtime.Autograd.TorchLean.adamw (Runtime.ofFloat lr) (Runtime.ofFloat weightDecay) (Runtime.ofFloat beta1) (Runtime.ofFloat beta2) (Runtime.ofFloat epsilon)).State) : have opt := Runtime.Autograd.TorchLean.adamw (Runtime.ofFloat lr) (Runtime.ofFloat weightDecay) (Runtime.ofFloat beta1) (Runtime.ofFloat beta2) (Runtime.ofFloat epsilon); IO (Runtime.Autograd.Train.DataLoader (TList α [σ, τ]) × opt.State)

structure NN.API.TorchLean.Supervised.Stepper (α : Type) [Semantics.Scalar α] [DecidableEq Spec.Shape] {σ τ : Spec.Shape} (task : SeqTask σ τ) : Type

Stateful training loop object: a Runner plus an optimizer state and a step counter.

This is the TorchLean analogue of holding a PyTorch optimizer object plus the model, ready to step() on batches.

• runner : Runner α task

  Underlying task runner (module + compiled predictors/losses).

• stepSample : TList α [σ, τ] → IO α

  Run a single optimization step on one supervised sample, returning the loss value.

• epochSamples : List (TList α [σ, τ]) → IO (List α)

  Run an epoch over an explicit list of samples, returning the per-step loss values.

• stepCount : IO ℕ

  Read the total number of stepSample calls performed so far.

def NN.API.TorchLean.Supervised.stepper {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] [Runtime.Scalar α] (runner : Runner α task) (optimizer : OptimizerConfig) (scheduler : Option Schedulers.Config := none) : IO (Stepper α task)

Construct a Stepper for a runner, optimizer config, and optional scheduler.

This is the recommended way to build custom training loops without reimplementing the optimizer logic: call stepper, then choose stepSample for single batches or epochSamples for explicit sample lists.

def NN.API.TorchLean.Supervised.step {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (loop : Stepper α task) (sample : TList α [σ, τ]) : IO α

Run one optimization step on a single supervised sample.

def NN.API.TorchLean.Supervised.epoch {σ τ : Spec.Shape} {task : SeqTask σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (loop : Stepper α task) (samples : List (TList α [σ, τ])) : IO (List α)

Run one epoch over a list of supervised samples, returning the per-step losses.

NN.API.TorchLean.Trainer is a stable facade over the internal Supervised training machinery.

Usage note:

If you're writing tutorials or user code, prefer the PyTorch-shaped entrypoint NN.API.train (available after import NN). This API.TorchLean.Trainer namespace is the underlying runtime layer that API.train delegates to, and it exposes more knobs than most users need.

The intended workflow is:

• pick a Task (regression / classification),
• call instantiate to get a Runner (parameters + buffers + backend state),
• call fit / fitDataset / fitLoader, or build a Stepper for custom loops.

This API is backend-agnostic: the same code can run in .eager mode or via a compiled backend, depending on the backend argument passed to instantiate.

@[reducible, inline]

abbrev NN.API.TorchLean.Trainer.Task (σ τ : Spec.Shape) : Type 2

A supervised task is just a model plus a choice of loss.

@[reducible, inline]

abbrev NN.API.TorchLean.Trainer.Runner (α : Type) [Semantics.Scalar α] [DecidableEq Spec.Shape] {σ τ : Spec.Shape} (task : Supervised.SeqTask σ τ) : Type

A fully instantiated supervised task runner.

This bundles:

• the imperative ScalarModule (parameters/buffers stored in refs),
• compiled predictors and loss functions for both .train and .eval modes (so switching mode is cheap),
• and the current mode stored in an IO.Ref.

The mode influences both operator behavior (e.g. dropout/batchnorm) and whether buffers are updated during training.

@[reducible, inline]

abbrev NN.API.TorchLean.Trainer.Optimizer : Type

Optimizer hyperparameter configuration for the supervised training helpers.

We keep this small for examples and lightweight trainers. It mirrors a few common PyTorch optimizers by name/defaults, but it does not try to cover the full option surface of torch.optim.*.

@[reducible, inline]

abbrev NN.API.TorchLean.Trainer.FitConfig : Type

Step-based training configuration for fit / fitDataset.

Fields:

• steps: number of parameter updates,
• optimizer: optimizer hyperparameters,
• scheduler: optional learning-rate schedule (applied per step),
• logEvery: progress printing frequency (0 disables logging).

@[reducible, inline]

abbrev NN.API.TorchLean.Trainer.LoaderFitConfig : Type

Epoch-based training configuration for fitLoader (data-loader training).

Fields:

• epochs: number of epochs (each epoch iterates once over the loader),
• optimizer: optimizer hyperparameters,
• scheduler: optional learning-rate schedule (applied per step/epoch depending on helper),
• logEvery: progress printing frequency (0 disables logging).

@[reducible, inline]

abbrev NN.API.TorchLean.Trainer.FitReport (α : Type) : Type

Small summary returned by fit* helpers.

By default, before and after are mean loss values, but the type is polymorphic so callers can report other scalars in the same shape.

@[reducible, inline]

abbrev NN.API.TorchLean.Trainer.Stepper (α : Type) [Semantics.Scalar α] [DecidableEq Spec.Shape] {σ τ : Spec.Shape} (task : Supervised.SeqTask σ τ) : Type

Stateful training loop object: a Runner plus an optimizer state and a step counter.

This is the TorchLean analogue of holding a PyTorch optimizer object plus the model, ready to step() on batches.

def NN.API.TorchLean.Trainer.regression {σ τ : Spec.Shape} (model : NN.Seq σ τ) (reduction : Loss.Reduction := Runtime.Autograd.TorchLean.Loss.Reduction.mean) : Task σ τ

Regression task with mean-squared error loss by default.

def NN.API.TorchLean.Trainer.classificationOneHot {σ τ : Spec.Shape} (model : NN.Seq σ τ) (reduction : Loss.Reduction := Runtime.Autograd.TorchLean.Loss.Reduction.mean) : Task σ τ

One-hot classification task with cross-entropy loss by default.

def NN.API.TorchLean.Trainer.sgd (lr : Float) (momentum : Float := 0.0) : Optimizer

SGD optimizer config.

PyTorch analogue: torch.optim.SGD (https://pytorch.org/docs/stable/generated/torch.optim.SGD.html).

def NN.API.TorchLean.Trainer.momentumSGD (lr : Float) (momentum : Float := 0.9) : Optimizer

Momentum SGD optimizer config (PyTorch-style default momentum = 0.9).

This is just sgd lr momentum with a different default.

def NN.API.TorchLean.Trainer.adam (lr : Float) (beta1 : Float := 0.9) (beta2 : Float := 0.999) (epsilon : Float := 1e-8) : Optimizer

Adam optimizer config with standard defaults.

PyTorch analogue: torch.optim.Adam (https://pytorch.org/docs/stable/generated/torch.optim.Adam.html).

def NN.API.TorchLean.Trainer.adamw (lr : Float) (weightDecay : Float := 1e-2) (beta1 : Float := 0.9) (beta2 : Float := 0.999) (epsilon : Float := 1e-8) : Optimizer

AdamW optimizer config with standard defaults (PyTorch-style weightDecay = 0.01).

PyTorch analogue: torch.optim.AdamW (https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html).

def NN.API.TorchLean.Trainer.steps (count : ℕ) (optimizer : Optimizer := sgd 1e-2) (logEvery : ℕ := 1) : FitConfig

Fixed-step training config over an in-memory sample list or dataset.

def NN.API.TorchLean.Trainer.epochs (count : ℕ) (optimizer : Optimizer := sgd 1e-2) (logEvery : ℕ := 1) : LoaderFitConfig

Epoch-based training config over a data loader.

def NN.API.TorchLean.Trainer.withScheduler (cfg : FitConfig) (scheduler : Schedulers.Config) : FitConfig

Attach a scheduler to a step-based training config.

def NN.API.TorchLean.Trainer.withEpochScheduler (cfg : LoaderFitConfig) (scheduler : Schedulers.Config) : LoaderFitConfig

Attach a scheduler to an epoch-based loader training config.

def NN.API.TorchLean.Trainer.constantLR (cfg : FitConfig) (lr : Float) : FitConfig

Step-based constant learning-rate schedule.

def NN.API.TorchLean.Trainer.stepLR (cfg : FitConfig) (base : Float) (stepSize : ℕ) (gamma : Float := 0.1) : FitConfig

Step-based step-decay schedule.

def NN.API.TorchLean.Trainer.exponentialLR (cfg : FitConfig) (base gamma : Float) : FitConfig

Step-based exponential-decay schedule.

def NN.API.TorchLean.Trainer.constantEpochLR (cfg : LoaderFitConfig) (lr : Float) : LoaderFitConfig

Epoch-based constant learning-rate schedule.

def NN.API.TorchLean.Trainer.stepEpochLR (cfg : LoaderFitConfig) (base : Float) (stepSize : ℕ) (gamma : Float := 0.1) : LoaderFitConfig

Epoch-based step-decay schedule.

def NN.API.TorchLean.Trainer.exponentialEpochLR (cfg : LoaderFitConfig) (base gamma : Float) : LoaderFitConfig

Epoch-based exponential-decay schedule.

def NN.API.TorchLean.Trainer.instantiateWithOptions {σ τ : Spec.Shape} (task : Task σ τ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [Runtime.Scalar α] [Runtime.Autograd.Torch.Internal.CudaBridge.TensorConv α] (opts : Options := { }) : IO (Runner α task)

Instantiate a runner under explicit Torch options (backend, useGpu, fastKernels, ...).

This is the recommended entrypoint when you want CUDA eager execution from the training helpers without dropping down to TorchLean.Module directly.

def NN.API.TorchLean.Trainer.instantiate {σ τ : Spec.Shape} (task : Task σ τ) {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [Runtime.Scalar α] [Runtime.Autograd.Torch.Internal.CudaBridge.TensorConv α] (backend : Backend := Runtime.Autograd.Torch.Backend.eager) : IO (Runner α task)

Instantiate a runner (parameters + buffers + backend state) for the given task.

This allocates and initializes model parameters (via Seq.initParams) and sets up the chosen execution backend (.eager vs .compiled).

def NN.API.TorchLean.Trainer.run {σ τ : Spec.Shape} (task : Task σ τ) (args : List String) (k : {α : Type} → [inst : Semantics.Scalar α] → [inst_1 : DecidableEq Spec.Shape] → [ToString α] → [Runtime.Scalar α] → Runner α task → List String → IO Unit) : IO Unit

CLI-oriented runner entry point.

This parses dtype/backend flags (via NN.API.DType / Module.ExecConfig) and then calls the continuation k under the selected scalar backend.

def NN.API.TorchLean.Trainer.params {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO (TList α (Supervised.paramShapes task))

Get the current model parameters from a runner.

def NN.API.TorchLean.Trainer.mode {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO NN.Mode

Read the current mode (train vs eval).

def NN.API.TorchLean.Trainer.setMode {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) (value : NN.Mode) : IO Unit

Set the mode (train vs eval).

def NN.API.TorchLean.Trainer.trainMode {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO Unit

Switch to training mode.

def NN.API.TorchLean.Trainer.evalMode {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO Unit

Switch to eval mode.

def NN.API.TorchLean.Trainer.isTraining {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) : IO Bool

Check whether the runner is in training mode.

def NN.API.TorchLean.Trainer.backward {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) (sample : TList α [σ, τ]) : IO (TList α (Supervised.paramShapes task))

Run forward+backward on one supervised sample and return gradients for all parameters.

def NN.API.TorchLean.Trainer.predict {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) (x : Spec.Tensor α σ) : IO (Spec.Tensor α τ)

Predict on a single input tensor.

This runs the forward pass under the runner's current mode.

def NN.API.TorchLean.Trainer.predictBatch {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (runner : Runner α task) (xs : List (Spec.Tensor α σ)) : IO (List (Spec.Tensor α τ))

Predict on a list of inputs (runs the forward pass repeatedly).

def NN.API.TorchLean.Trainer.predictClass? {σ : Spec.Shape} {n : ℕ} {task : Task σ (Spec.Shape.dim n Spec.Shape.scalar)} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (runner : Runner α task) (x : Spec.Tensor α σ) : IO (Option (Fin n))

Predict the argmax class for a classification task, if argmax is well-defined for α.

This is a convenience wrapper over predict + Metrics.argmax?.

def NN.API.TorchLean.Trainer.accuracyOneHot {σ : Spec.Shape} {n : ℕ} {task : Task σ (Spec.Shape.dim n Spec.Shape.scalar)} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [LT α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (runner : Runner α task) (samples : List (TList α [σ, Spec.Shape.dim n Spec.Shape.scalar])) : IO (ℕ × ℕ)

Count correct predictions in a one-hot labeled sample list (returns (correct, total)).

def NN.API.TorchLean.Trainer.meanLoss {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] (runner : Runner α task) (samples : List (TList α [σ, τ])) : IO α

Mean loss over an explicit list of samples.

def NN.API.TorchLean.Trainer.meanLossDataset {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] (runner : Runner α task) (dataset : Runtime.Autograd.Train.Dataset (TList α [σ, τ])) : IO α

Mean loss over an entire Dataset.

def NN.API.TorchLean.Trainer.fit {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] [Runtime.Scalar α] (runner : Runner α task) (cfg : FitConfig) (samples : List (TList α [σ, τ])) : IO (FitReport α)

Fit on an explicit list of samples for a fixed number of steps.

Returns a small report with mean loss before/after.

def NN.API.TorchLean.Trainer.fitDataset {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] [Runtime.Scalar α] (runner : Runner α task) (cfg : FitConfig) (dataset : Runtime.Autograd.Train.Dataset (TList α [σ, τ])) : IO (FitReport α)

Fit on a Dataset for a fixed number of steps.

def NN.API.TorchLean.Trainer.fitLoader {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] [Runtime.Scalar α] (runner : Runner α task) (cfg : LoaderFitConfig) (loader : Runtime.Autograd.Train.DataLoader (TList α [σ, τ])) : IO (FitReport α × Runtime.Autograd.Train.DataLoader (TList α [σ, τ]))

Fit using a DataLoader for a fixed number of epochs.

def NN.API.TorchLean.Trainer.stepper {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] [Runtime.Scalar α] (runner : Runner α task) (optimizer : Optimizer) (scheduler : Option Schedulers.Config := none) : IO (Stepper α task)

Construct a stateful stepper for custom loops.

This is useful if you want to control:

• evaluation cadence,
• logging,
• validation, early stopping, etc.

The returned Stepper still uses TorchLean's optimizer/scheduler implementations.

def NN.API.TorchLean.Trainer.step {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (loop : Stepper α task) (sample : TList α [σ, τ]) : IO α

Run a single training step on one sample using a Stepper.

def NN.API.TorchLean.Trainer.epoch {σ τ : Spec.Shape} {task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] (loop : Stepper α task) (samples : List (TList α [σ, τ])) : IO (List α)

Run an epoch over a list of samples using a Stepper (returns the per-step losses).

def NN.API.TorchLean.Trainer.train {σ τ : Spec.Shape} {_task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] [Runtime.Scalar α] (task : Task σ τ) (cfg : FitConfig) (samples : List (TList α [σ, τ])) (backend : Backend := Runtime.Autograd.Torch.Backend.eager) : IO (Runner α task × FitReport α)

Convenience: instantiate + fit on a list of samples.

Returns both the Runner (so you can keep using the trained parameters) and the fit report.

def NN.API.TorchLean.Trainer.trainDataset {σ τ : Spec.Shape} {_task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] [Runtime.Scalar α] (task : Task σ τ) (cfg : FitConfig) (dataset : Runtime.Autograd.Train.Dataset (TList α [σ, τ])) (backend : Backend := Runtime.Autograd.Torch.Backend.eager) : IO (Runner α task × FitReport α)

Convenience: instantiate + fit on a Dataset.

Returns both the Runner and the fit report.

def NN.API.TorchLean.Trainer.trainLoader {σ τ : Spec.Shape} {_task : Task σ τ} {α : Type} [Semantics.Scalar α] [DecidableEq Spec.Shape] [ToString α] [Add α] [Div α] [Zero α] [Coe ℕ α] [Runtime.Scalar α] (task : Task σ τ) (cfg : LoaderFitConfig) (loader : Runtime.Autograd.Train.DataLoader (TList α [σ, τ])) (backend : Backend := Runtime.Autograd.Torch.Backend.eager) : IO (Runner α task × FitReport α × Runtime.Autograd.Train.DataLoader (TList α [σ, τ]))

Convenience: instantiate + fit using a DataLoader.

Returns the Runner, the fit report, and the updated loader state (shuffled epoch cursor).

@[reducible, inline]

abbrev NN.API.TorchLean.ExecConfig : Type

Execution config parsed from CLI flags (dtype/backend/fast-kernels).

@[reducible, inline]

abbrev NN.API.TorchLean.ExecConfig.parseAndStrip (args : List String) : Except String (Module.ExecConfig × List String)

Parse and strip execution flags, returning (config, remainingArgs).

@[reducible, inline]

abbrev NN.API.TorchLean.ExecConfig.log (cfg : Module.ExecConfig) : IO Unit

Log the chosen execution config to stdout (useful for reproducible demos).

Re-export of the low-level imperative scalar trainer interface.

This exposes forwardT/backwardT/stepT from Runtime.Autograd.TorchLean.ScalarTrainer. Most users should prefer the higher-level NN.API.train / API.TorchLean.Trainer APIs.

Imperative session API: a tape-backed interface that can run in eager or compiled mode.

This is roughly analogous to using PyTorch "eager tensors", except TorchLean makes the tape/session explicit. The Session surface is useful for:

• interactive experiments in IO,
• debugging (inspect intermediate values),
• building higher-level runners.