NN

TorchLean.NN: a compact torch.nn-style builder layer.

This module defines a small torch.nn-style builder layer for constructing shape-typed models. It packages model-state shapes and initial values together with an execution-polymorphic forward program, so example code does not have to spell paramShapes := [...] / inputShapes := [...] everywhere.

Main definitions

• Layer σ τ packages a shape-typed layer with explicit state (parameters and buffers) and a polymorphic forward program.
• Seq σ τ composes layers sequentially (PyTorch analogy: torch.nn.Sequential), written f >>> g.
• Seq.Objective bundles a Seq model together with a scalar loss, producing a TorchLean.Module.ObjectiveDef that the runtime training code can execute.

PyTorch analogies

• Layer is like a small nn.Module definition, except parameters are an explicit list instead of fields, and the forward pass is a typed TorchLean program.
• Mode is like module.train() vs module.eval() (dropout and batchnorm-like layers branch on it).
• The updateBuffers mechanism is like updating non-gradient buffers (e.g. BatchNorm running stats).

The surface here is narrow by design: it supports TorchLean's executable model constructors and training helpers without trying to mirror the full torch.nn API.

References

• PyTorch torch.nn: https://pytorch.org/docs/stable/nn.html

Mode

TorchLean keeps "train vs eval" behavior explicit. This affects layers like dropout and batch-normalization that behave differently during training vs inference.

inductive Runtime.Autograd.Model.Layers.Mode : Type

Execution mode for layers that branch between training-time and inference-time behavior.

PyTorch analogy: model.train() / model.eval() (affects dropout, batchnorm, etc.).

• train : Mode
• eval : Mode

@[instance_reducible]

instance Runtime.Autograd.Model.Layers.instReprMode : Repr Mode

def Runtime.Autograd.Model.Layers.instReprMode.repr : Mode → ℕ → Std.Format

@[instance_reducible]

instance Runtime.Autograd.Model.Layers.instDecidableEqMode : DecidableEq Mode

Layer definitions

structure Runtime.Autograd.Model.Layers.Layer (σ τ : Spec.Shape) : Type 1

A shape-typed layer definition with explicit model state and an execution-polymorphic forward program.

Layer σ τ is the core building block used by Seq (sequential composition). It stores:

• the shapes of parameters and persistent buffers,
• initial values for that state (as Float tensors, for reproducible initialization),
• per-parameter requiresGrad flags, and
• a forward program that is polymorphic over the backend monad and scalar type.

PyTorch analogy: a small nn.Module, where:

• stateShapes/initState contain parameters and persistent buffers,
• forward corresponds to Module.forward,
• updateBuffers corresponds to updating things like running_mean/running_var in BatchNorm.

• kind : String

  Layer label used by public model summaries.

• stateShapes : List Spec.Shape

  Shapes of parameters and persistent buffers, in the order expected by forward.

• initState : TorchLean.TensorPack Float self.stateShapes

  Initial model state, stored as Float tensors for convenient initialization.

• runtimeInit : Option (Module.RuntimeInit.Plan self.stateShapes)

  Optional storage-first initialization plan for executable Float backends.

  This does not replace initState: the tensor-valued initializers remain available to the specification and proof layers. The plan lets a runtime create equivalent parameter storage without first enumerating those tensors on the host.

• requiresGrad : Array Bool

  Gradient flags for model state (defaults to all true). Buffers use false.

  PyTorch analogy: tensor.requires_grad_(...) on parameters/buffers.

• validateConfig : Except String Unit

  Validate static layer configuration before allocating runtime state or lowering a graph.

  Shape compatibility remains enforced by the type. This check is for value-level configuration such as a dropout probability that must belong to a finite numeric interval.

• updateBuffers : Option (Mode → {α : Type} → [inst : TorchLean.Storage α] → [Context α] → TorchLean.TensorPack α self.stateShapes → TorchLean.Tensor α σ → IO (TorchLean.TensorPack α self.stateShapes))

  Optional buffer update function (used for running-statistics style layers).

  This is called during a forward pass (typically in Mode.train) to produce updated parameter/buffer state values. A canonical example is BatchNorm updating its running_mean / running_var buffers.

• updatesBuffersInForward : Bool

  Composite layers delegate runtime buffer updates to their nested forward programs.

• forward : Mode → {α : Type} → [inst : TorchLean.Storage α] → [inst✝ : Context α] → Program α (self.stateShapes ++ [σ]) τ

  Forward pass as a typed TorchLean program.

  The program expects (stateShapes ++ [σ]) inputs (model state, then the layer input) and produces an output of shape τ.

def Runtime.Autograd.Model.Layers.updateRunning {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (running batch : TorchLean.Tensor α s) (momentum : TorchLean.Tensor α Spec.Shape.scalar) : TorchLean.Tensor α s

Update running statistics of any shape using momentum.

This implements an exponential moving average:

next = (1 - momentum) * running + momentum * batch.

PyTorch analogy: the update performed for running_mean / running_var in BatchNorm.

def Runtime.Autograd.Model.Layers.unbiasedRunningVariance {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (biased : TorchLean.Tensor α s) (sampleCount : ℕ) : TorchLean.Tensor α s

Convert the biased variance used by BatchNorm's training forward pass into the unbiased estimate stored in its running buffer. For a singleton sample set there is no unbiased estimate; TorchLean keeps the finite biased value rather than dividing by zero.

def Runtime.Autograd.Model.Layers.batchChannelStats {α : Type} [TorchLean.Storage α] [Context α] {batch channels : ℕ} {spatial : Spec.Shape} (x : TorchLean.Tensor α (Spec.Shape.dim batch (Spec.Shape.dim channels spatial))) : TorchLean.Tensor α [channels] × TorchLean.Tensor α [channels]

Compute per-channel mean and biased variance for a batched tensor.

The first two axes are batch and channel; every remaining axis is reduced. The result is a pair of vectors indexed by channel. A running-variance update uses unbiasedRunningVariance vars (batch * spatial.size) instead.

def Runtime.Autograd.Model.Layers.Layer.validate {σ τ : Spec.Shape} (layer : Layer σ τ) : Except String Unit

Validate a layer's complete static contract.

Configuration checks supplied by the layer are combined with generic state-metadata and runtime initializer checks so every execution path rejects malformed custom layers consistently.

def Runtime.Autograd.Model.Layers.Layer.forwardRef {σ τ : Spec.Shape} (l : Layer σ τ) {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] (mode : Mode) (ps : RefList (RefTy m α) l.stateShapes) (x : RefTy m α σ) : m (RefTy m α τ)

Run a Layer forward given parameter refs and an input ref.

This is the "module forward" operation at the reference level.

PyTorch analogy: calling layer(x) where the layer's parameters are already allocated.

def Runtime.Autograd.Model.Layers.Layer.forwardTensor {σ τ : Spec.Shape} (l : Layer σ τ) (mode : Mode) {α : Type} [TorchLean.Storage α] [Context α] (ps : TorchLean.TensorPack α l.stateShapes) (x : TorchLean.Tensor α σ) : IO (TorchLean.Tensor α τ)

Run a Layer on concrete tensors by lowering its forward program to a typed graph.

This is primarily used by runtime utilities (e.g. sequential updateBuffers) where we want to run forward to obtain intermediate activations.

PyTorch analogy: running a forward pass eagerly on concrete tensors.