Seeded model builders

Every public layer constructor lives here and returns nn.Builder, drawing deterministic initialization seeds from an explicit seed stream. The explicit-seed implementations behind these builders live in NN.API.Neural.Impl, which only this module imports; nn.build seed is the way to turn a builder into a model with fixed seeds.

nn.Sequential lives in Type 1 because every layer stores an execution-polymorphic forward program, so it cannot be returned directly from IO. We draw a base seed in IO, then use nn.build to construct the model purely. nn.buildIO returns the drawn seed in a Built value that IO can carry, and nn.withModel passes the resulting model to a continuation.

@[reducible, inline]

abbrev TorchLean.nn.Builder (α : Type u_1) : Type u_1

Deterministic model builder that threads an explicit initialization seed stream.

def TorchLean.nn.build {α : Type u} (seed : ℕ) (builder : Builder α) : α

Build a value from a deterministic initialization seed.

Example:

def model : nn.Builder (nn.Sequential [2] [1]) :=
  nn.Sequential![nn.linear 2 8, nn.relu, nn.linear 8 1]

-- Same seed, same weights, on every machine and every run.
def built : nn.Sequential [2] [1] := nn.build 7 model

def TorchLean.nn.withSeed {α : Type u} (continuation : ℕ → Builder α) : Builder α

Consume one initialization seed and continue building in the same result universe.

def TorchLean.nn.withInitializationSeed {α : Type u} (initialization : Init.Scheme) (continuation : ℕ → Builder α) : Builder α

Consume a seed exactly when an initialization scheme is stochastic.

def TorchLean.nn.withDropoutSeed {α : Type u} (probability : Float) (continuation : ℕ → Builder α) : Builder α

Consume a seed exactly when training dropout requires a random mask.

def TorchLean.nn.withOptionalDropoutSeed {α : Type u} (probability? : Option Float) (continuation : ℕ → Builder α) : Builder α

An absent dropout site consumes no key, just like a deterministic endpoint probability.

Layer constructors

Each builder validates its configuration, consumes exactly the seeds its stochastic initializers need, and delegates to nn.Impl.

def TorchLean.nn.globalAvgPool {d channels : ℕ} (spatial : Tensor ℕ [d]) (batchShape : Shape := []) : Builder (Sequential (batchShape.concat ((spatial.to Shape).prependDim channels)) (batchShape.appendDim channels))

Build global average pooling without consuming an initialization seed.

Example:

-- Average every `8 x 8` map down to one number per channel, the usual last step before a
-- classifier head.
def model : nn.Builder (nn.Sequential [3, 8, 8] [3]) :=
  nn.globalAvgPool [8, 8] (channels := 3)

def TorchLean.nn.heads.classifier {featureShape : Shape} (classCount : ℕ) (batchShape : Shape := []) : Builder (Sequential (batchShape.concat featureShape) (batchShape.appendDim classCount))

Build a classification head that flattens the feature suffix and seeds its weight.

def TorchLean.nn.heads.regressor {featureShape : Shape} (outputWidth : ℕ := 1) (batchShape : Shape := []) : Builder (Sequential (batchShape.concat featureShape) (batchShape.appendDim outputWidth))

Build a regression head that flattens the feature suffix and seeds its weight.

def TorchLean.nn.relu {shape : Shape} : Builder (Sequential shape shape)

Build an elementwise ReLU layer without consuming an initialization seed.

def TorchLean.nn.silu {shape : Shape} : Builder (Sequential shape shape)

Build an elementwise SiLU layer without consuming an initialization seed.

def TorchLean.nn.gelu {shape : Shape} : Builder (Sequential shape shape)

Build tanh-approximate GELU without consuming an initialization seed.

This retains the historical nn.gelu behavior. geluTanh names the same formula explicitly. Erf-based GELU is not supported by the scalar operation interface.

def TorchLean.nn.geluTanh {shape : Shape} : Builder (Sequential shape shape)

Explicit constructor for GELU's cubic tanh approximation.

def TorchLean.nn.sigmoid {shape : Shape} : Builder (Sequential shape shape)

Build an elementwise sigmoid layer without consuming an initialization seed.

def TorchLean.nn.tanh {shape : Shape} : Builder (Sequential shape shape)

Build an elementwise hyperbolic-tangent layer without consuming an initialization seed.

def TorchLean.nn.softmax {shape : Shape} (axis : ℕ) : Builder (Sequential shape shape)

Build a softmax layer along any valid tensor dimension without consuming a seed.

Example:

-- Axis `0` of a rank-one shape: a probability vector over ten classes.
def model : nn.Builder (nn.Sequential [10] [10]) :=
  nn.softmax 0

def TorchLean.nn.logSoftmax {shape : Shape} (axis : ℕ) : Builder (Sequential shape shape)

Build a stable log-softmax layer along any valid tensor dimension without consuming a seed.

def TorchLean.nn.sum {shape : Shape} : Builder (Sequential shape [])

Build a reduction that sums every tensor entry to a scalar.

def TorchLean.nn.flatten {shape : Shape} : Builder (Sequential shape [shape.size])

Build a layer that flattens the entire input shape into one vector.

def TorchLean.nn.reshape (source target : Shape) : Builder (Sequential source target)

Build a reshape that is rejected by model validation when the element counts differ.

def TorchLean.nn.flattenAfter (batchShape : Shape := []) {shape : Shape} : Builder (Sequential (batchShape.concat shape) (batchShape.appendDim shape.size))

Flatten each tensor after an arbitrary batch shape.

def TorchLean.nn.maxPool {d channels : ℕ} (spatial : Tensor ℕ [d]) (config : Pooling.Config d) (batchShape : Shape := []) : Builder (Sequential (batchShape.concat ((spatial.to Shape).prependDim channels)) (batchShape.concat (((config.outputSpatial spatial).to Shape).prependDim channels)))

Build max pooling over arbitrary spatial rank using the supplied pooling configuration.

Example:

def spatial : Tensor Nat [2] := [8, 8]

def pooling : nn.Pooling.Config 2 :=
  { kernelSize := [2, 2], stride := [2, 2] }

-- Non-overlapping `2 x 2` windows halve both spatial axes and leave the channel count alone.
def model : nn.Builder (nn.Sequential [3, 8, 8] [3, 4, 4]) :=
  nn.maxPool spatial pooling (channels := 3)

def TorchLean.nn.avgPool {d channels : ℕ} (spatial : Tensor ℕ [d]) (config : Pooling.Config d) (batchShape : Shape := []) : Builder (Sequential (batchShape.concat ((spatial.to Shape).prependDim channels)) (batchShape.concat (((config.outputSpatial spatial).to Shape).prependDim channels)))

Build average pooling over arbitrary spatial rank using the supplied pooling configuration.

def TorchLean.nn.convTranspose {d inputChannels : ℕ} (spatial : Tensor ℕ [d]) (config : TransposedConvolution.Config d) (batchShape : Shape := []) : Builder (Sequential (batchShape.concat ((spatial.to Shape).prependDim inputChannels)) (batchShape.concat (((config.outputSpatial spatial).to Shape).prependDim config.outChannels)))

Build a transpose convolution over arbitrary spatial rank.

def TorchLean.nn.linear (inputWidth outputWidth : ℕ) (batchShape : Shape := []) (config : Linear.Config := { }) : Builder (Sequential (batchShape.appendDim inputWidth) (batchShape.appendDim outputWidth))

Build an affine layer, consuming seeds only for stochastic initializers.

Example:

-- Two affine layers around a ReLU: `2 -> 8 -> 1`.
def model : nn.Builder (nn.Sequential [2] [1]) :=
  nn.Sequential![
    nn.linear 2 8,
    nn.relu,
    nn.linear 8 1
  ]

-- With a batch axis in front, the same call reads `[16, 2] -> [16, 8]`.
def batched : nn.Builder (nn.Sequential [16, 2] [16, 8]) :=
  nn.linear 2 8 (batchShape := [16])

def TorchLean.nn.rnn (sequenceLength inputWidth hiddenWidth : ℕ) (batchShape : Shape := []) : Builder (Sequential ((batchShape.appendDim sequenceLength).appendDim inputWidth) ((batchShape.appendDim sequenceLength).appendDim hiddenWidth))

Build a seeded recurrent neural network over a fixed sequence length.

The recurrent computation acts independently over every index in batchShape; its parameters are shared across those indices. The scalar default is a single sequence, while a shape such as [batch] gives the usual batched model.

def TorchLean.nn.gru (sequenceLength inputWidth hiddenWidth : ℕ) (batchShape : Shape := []) (convention : Spec.GRUConvention := Spec.GRUConvention.resetBefore) : Builder (Sequential ((batchShape.appendDim sequenceLength).appendDim inputWidth) ((batchShape.appendDim sequenceLength).appendDim hiddenWidth))

Build a seeded GRU, shared over every index in batchShape.

resetBefore preserves the original constructor and its three initialization draws. Choose resetAfter for PyTorch's recurrence and four packed parameter tensors; that version draws one seed for each of its input and recurrent weight matrices and keeps both biases independent.

def TorchLean.nn.gruFromPyTorch (sequenceLength : ℕ) {inputWidth hiddenWidth : ℕ} (parameters : Spec.GRUResetAfterSpec Float inputWidth hiddenWidth) (batchShape : Shape := []) : Builder (Sequential ((batchShape.appendDim sequenceLength).appendDim inputWidth) ((batchShape.appendDim sequenceLength).appendDim hiddenWidth))

Load one reset-after GRU cell's parameters and share it over every batch position.

Pass Spec.GRUResetAfterSpec.ofPyTorch weightIH weightHH biasIH biasHH. The gate order and matrix layout are retained, and no initialization seeds are consumed.

def TorchLean.nn.mamba (sequenceLength inputWidth hiddenWidth : ℕ) (batchShape : Shape := []) (options : Runtime.Autograd.Model.Mamba.Options := { }) : Builder (Sequential ((batchShape.appendDim sequenceLength).appendDim inputWidth) ((batchShape.appendDim sequenceLength).appendDim hiddenWidth))

Build a seeded selective Mamba layer, shared over every index in batchShape.

def TorchLean.nn.lstm (sequenceLength inputWidth hiddenWidth : ℕ) (batchShape : Shape := []) : Builder (Sequential ((batchShape.appendDim sequenceLength).appendDim inputWidth) ((batchShape.appendDim sequenceLength).appendDim hiddenWidth))

Build a seeded long short-term memory layer, shared over every index in batchShape.

Example:

-- Sixteen timesteps of width 8 in, sixteen hidden states of width 32 out.
def model : nn.Builder (nn.Sequential [16, 8] [16, 32]) :=
  nn.lstm 16 8 32

def TorchLean.nn.conv {d inputChannels : ℕ} (spatial : Tensor ℕ [d]) (config : Convolution.Config d) (batchShape : Shape := []) : Builder (Sequential (batchShape.concat ((spatial.to Shape).prependDim inputChannels)) (batchShape.concat (((config.outputSpatial spatial).to Shape).prependDim config.outChannels)))

Build an arbitrary-rank convolution, seeding its kernel when initialization is stochastic.

Example:

def spatial : Tensor Nat [2] := [8, 8]

def convolution : nn.Convolution.Config 2 :=
  { outChannels := 4, kernelSize := [3, 3] }

-- One `8 x 8` channel in, four `6 x 6` feature maps out: no padding, so the kernel eats a
-- one-pixel border on each side.
def model : nn.Builder (nn.Sequential [1, 8, 8] [4, 6, 6]) :=
  nn.conv spatial convolution (inputChannels := 1)

def TorchLean.nn.pointwiseConv {d inputChannels : ℕ} (spatial : Tensor ℕ [d]) (outputChannels : ℕ) (batchShape : Shape := []) : Builder (Sequential (batchShape.concat ((spatial.to Shape).prependDim inputChannels)) (batchShape.concat ((spatial.to Shape).prependDim outputChannels)))

Build a pointwise convolution over any spatial rank.

The unit kernel, unit stride, and zero padding preserve every spatial axis. This is the common channel-projection operation used by residual, diffusion, and encoder-decoder models.

def TorchLean.nn.batchNorm {d channels : ℕ} (spatial : Tensor ℕ [d]) (momentum : Float := 0.1) (batchShape : Shape := []) (eps : ℚ := 1e-5) : Builder (Sequential (batchShape.concat ((spatial.to Shape).prependDim channels)) (batchShape.concat ((spatial.to Shape).prependDim channels)))

Build batch normalization with learned scale and bias and stored running statistics.

Training computes one mean and variance per channel across the batch and spatial axes. Evaluation uses the stored statistics. momentum controls how much each training batch changes those buffers; eps is added to the variance before taking its square root. Scale and running variance start at one, while bias and running mean start at zero.

Example:

-- Shape in equals shape out. What changes is the running mean and variance this layer keeps as
-- persistent buffers, updated in `.train` mode and only read in `.eval` mode.
def model : nn.Builder (nn.Sequential [3, 8, 8] [3, 8, 8]) :=
  nn.batchNorm [8, 8] (momentum := 0.1) (channels := 3)

def TorchLean.nn.instanceNorm {d channels : ℕ} (spatial : Tensor ℕ [d]) (batchShape : Shape := []) (eps : ℚ := 1e-5) (affine bias : Bool := true) : Builder (Sequential (batchShape.concat ((spatial.to Shape).prependDim channels)) (batchShape.concat ((spatial.to Shape).prependDim channels)))

Build instance normalization using each sample and channel's spatial mean and variance.

The layer uses current input statistics in both training and evaluation. Scale and bias each have one entry per channel and start at one and zero. Both are enabled by default; bias := false keeps only the scale, and affine := false removes both. eps is added to the variance before taking its square root.

def TorchLean.nn.groupNorm {d channels : ℕ} (spatial : Tensor ℕ [d]) (groups : ℕ) (batchShape : Shape := []) (eps : ℚ := 1e-5) (affine bias : Bool := true) : Builder (Sequential (batchShape.concat ((spatial.to Shape).prependDim channels)) (batchShape.concat ((spatial.to Shape).prependDim channels)))

Build group normalization with equal, contiguous groups of channels.

Each sample's groups have separate means and variances, computed over their channels and spatial positions. The channel count must be divisible by the positive group count. Scale and bias have one entry per channel; bias := false keeps only the scale, and affine := false removes both. eps is added to each group's variance before taking its square root.

def TorchLean.nn.oneHotEmbedding (vocabularySize embeddingWidth : ℕ) (config : Embedding.Config := { }) (batchShape : Shape := []) : Builder (Sequential (batchShape.appendDim vocabularySize) (batchShape.appendDim embeddingWidth))

Build an embedding lookup layer from a freshly seeded embedding table.

def TorchLean.nn.embedding (vocabularySize embeddingWidth : ℕ) (config : Embedding.Config := { }) : Builder (Embedding vocabularySize embeddingWidth)

Build a trainable lookup table for a tensor of natural-number indices.

Example:

-- A vocabulary of 256 byte tokens, each mapped to a 32-dimensional row.
def table : nn.Embedding 256 32 :=
  nn.build 0 (nn.embedding 256 32)

-- `table.model` fixes the index shape: a length-16 token window becomes `[16, 32]`.
def model : nn.IndexedModel [16] [16, 32] (Fin 256) :=
  table.model [16]

def TorchLean.nn.sinusoidalPositionalEncoding (batchShape : Shape := []) {sequenceLength embeddingWidth : ℕ} (config : SinusoidalPositionalEncoding.Config := { }) : Builder (Sequential (batchShape.concat [sequenceLength, embeddingWidth]) (batchShape.concat [sequenceLength, embeddingWidth]))

Build deterministic sinusoidal positional encoding over a sequence suffix.

def TorchLean.nn.rope (batchShape : Shape := []) {sequenceLength headWidth : ℕ} (config : RotaryEmbedding.Config := { }) : Builder (Sequential (batchShape.concat [sequenceLength, headWidth]) (batchShape.concat [sequenceLength, headWidth]))

Build deterministic rotary positional encoding for multi-head sequence features.

def TorchLean.nn.learnedPositionalEmbedding (batchShape : Shape := []) {sequenceLength embeddingWidth : ℕ} (config : LearnedPositionalEmbedding.Config := { }) : Builder (Sequential (batchShape.concat [sequenceLength, embeddingWidth]) (batchShape.concat [sequenceLength, embeddingWidth]))

Build learned positional embeddings from a freshly allocated parameter seed.

def TorchLean.nn.layerNorm (batchShape : Shape := []) {width : ℕ} (eps : ℚ := 1e-5) (affine bias : Bool := true) : Builder (Sequential (batchShape.appendDim width) (batchShape.appendDim width))

Build layer normalization over the final axis.

Each row uses its own mean and variance. Scale and bias have shape [width] and are shared across the leading axes, starting at one and zero. eps is added to the variance before taking its square root. Setting bias := false keeps only the scale; setting affine := false removes both parameters.

The rational eps must remain positive and finite after conversion to the execution scalar. The default can round to zero in tiny formats, producing NaNs on constant rows in a typed graph. Validation checks rational positivity only. Choose a representable positive value with eps; for three exponent bits and two fraction bits, (eps := (1 / 16 : Rat)) is such a value.

Example:

-- Normalizes across the final axis of each `[16, 64]` row, the Transformer convention.
def model : nn.Builder (nn.Sequential [16, 64] [16, 64]) :=
  nn.layerNorm [16] (width := 64)

def TorchLean.nn.rmsNorm (batchShape : Shape := []) {width : ℕ} (eps : ℚ := 1e-5) (affine : Bool := true) : Builder (Sequential (batchShape.appendDim width) (batchShape.appendDim width))

Build RMS normalization over the final axis.

Each row is divided by sqrt(mean(x * x) + eps), then multiplied by a scale of shape [width]. There is no mean subtraction or bias. The scale starts at one; affine := false removes it from model state. The default eps is 1e-5 for every scalar type. The converted epsilon must remain positive and finite. If it rounds to zero in a tiny format, zero input can produce NaNs; pass a representable positive eps before lowering the model.

def TorchLean.nn.multiHeadAttention {sequenceLength modelWidth : ℕ} (config : MultiHeadAttention.Config) (mask : Option (Tensor Bool [sequenceLength, sequenceLength]) := none) (batchShape : Shape := []) : Builder (Sequential (batchShape.concat [sequenceLength, modelWidth]) (batchShape.concat [sequenceLength, modelWidth]))

Build seeded multi-head self-attention with an optional fixed attention mask.

Example:

-- Two heads of width 4 give an internal attention width of 8, which here happens to match the
-- model width; the two are independent, so `headCount * headWidth` may differ from it.
def model : nn.Builder (nn.Sequential [4, 8] [4, 8]) :=
  nn.multiHeadAttention { headCount := 2, headWidth := 4 }
    (sequenceLength := 4) (modelWidth := 8)

-- Causal masking is a separate argument rather than a config field, because the mask is a value
-- with the sequence length in its type.
def causal : nn.Builder (nn.Sequential [4, 8] [4, 8]) :=
  nn.multiHeadAttention { headCount := 2, headWidth := 4 }
    (mask := some (Spec.causalMask 4)) (sequenceLength := 4) (modelWidth := 8)

def TorchLean.nn.transformerEncoderBlock {sequenceLength modelWidth : ℕ} (config : TransformerEncoder.Block.Config) (mask : Option (Tensor Bool [sequenceLength, sequenceLength]) := none) (batchShape : Shape := []) : Builder (Sequential (batchShape.concat [sequenceLength, modelWidth]) (batchShape.concat [sequenceLength, modelWidth]))

Build one seeded transformer encoder block, optionally applying a fixed attention mask.

Example:

-- Pre-norm block, GELU feed-forward, no dropout: attention and feed-forward each sit inside their
-- own residual connection, so shapes in and out agree.
def model : nn.Builder (nn.Sequential [16, 64] [16, 64]) :=
  nn.transformerEncoderBlock
    { headCount := 4
      headWidth := 16
      feedForwardWidth := 256
      normalizeFirst := true }
    (sequenceLength := 16) (modelWidth := 64)

def TorchLean.nn.transformerEncoderStack {sequenceLength modelWidth : ℕ} (config : TransformerEncoder.Stack.Config) (mask : Option (Tensor Bool [sequenceLength, sequenceLength]) := none) (batchShape : Shape := []) : Builder (Sequential (batchShape.concat [sequenceLength, modelWidth]) (batchShape.concat [sequenceLength, modelWidth]))

Build a seeded stack of transformer encoder blocks with an optional attention mask.

def TorchLean.nn.transformerEncoderStack.buildLayers {sequenceLength modelWidth : ℕ} (config : TransformerEncoder.Stack.Config) (mask : Option (Tensor Bool [sequenceLength, sequenceLength])) (batchShape : Shape) : ℕ → have shape := batchShape.concat [sequenceLength, modelWidth]; Builder (Sequential shape shape)

def TorchLean.nn.dropout {shape : Shape} (p : Float) : Builder (Sequential shape shape)

Build dropout, consuming a key only when a random training mask is possible.

Example:

-- Active in `.train` mode and the identity in `.eval` mode, which the trainer selects for you.
def model : nn.Builder (nn.Sequential [64] [64]) :=
  nn.dropout 0.1

def TorchLean.nn.convBlock {d inputChannels : ℕ} (spatial : Tensor ℕ [d]) (config : ConvBlock.Config d) (batchShape : Shape := []) : Builder (Sequential (batchShape.concat ((spatial.to Shape).prependDim inputChannels)) (batchShape.concat (((config.convolution.outputSpatial spatial).to Shape).prependDim config.convolution.outChannels)))

Build a seeded convolution, activation, and optional dropout block.

def TorchLean.nn.convPoolBlock {d inputChannels : ℕ} (spatial : Tensor ℕ [d]) (config : ConvPoolBlock.Config d) (batchShape : Shape := []) : Builder (Sequential (batchShape.concat ((spatial.to Shape).prependDim inputChannels)) (batchShape.concat (((config.pooling.outputSpatial (config.block.convolution.outputSpatial spatial)).to Shape).prependDim config.block.convolution.outChannels)))

Build a seeded convolution/activation block followed by max pooling.

def TorchLean.nn.mlp (inputWidth outputWidth : ℕ) (config : MLP.Config := { }) (batchShape : Shape := []) : Builder (Sequential (batchShape.appendDim inputWidth) (batchShape.appendDim outputWidth))

Build a multilayer perceptron over any batchShape.

Each hidden width contributes a linear layer followed by the configured activation and optional dropout. Initialization seeds come from the surrounding Builder seed stream.

Example:

-- `16 -> 32 -> 32 -> 1`, ReLU between hidden layers, dropout after each one.
def model : nn.Builder (nn.Sequential [16] [1]) :=
  nn.mlp 16 1 { hiddenWidths := [32, 32], activation := .relu, dropout? := some 0.1 }

def TorchLean.nn.mlp.buildStages (outputWidth : ℕ) (config : MLP.Config) (batchShape : Shape) (currentWidth : ℕ) (hiddenWidths : List ℕ) : Builder (Sequential (batchShape.appendDim currentWidth) (batchShape.appendDim outputWidth))

structure TorchLean.nn.Built {σ τ : Shape} (builder : Builder (Sequential σ τ)) : Type

A model built from builder with a seed drawn at run time.

nn.Sequential lives in Type 1, so IO cannot return it directly. This record stores only the drawn seed; because builder is deterministic, Built.model recovers the same model every time. Built builder lives in Type, so it can be returned from IO and stored in ordinary records.

• seed : ℕ

  The initialization seed drawn for this model.

def TorchLean.nn.Built.model {σ τ : Shape} {builder : Builder (Sequential σ τ)} (built : Built builder) : Sequential σ τ

The model determined by a Built value.

def TorchLean.nn.buildIO {σ τ : Shape} (builder : Builder (Sequential σ τ)) : IO (Built builder)

Draw the next global seed for builder and return it as an IO value.

This is the direct-construction counterpart of nn.withModel: (← nn.buildIO builder).model is the model that nn.withModel builder would pass to its continuation.

def TorchLean.nn.withModel {σ τ : Shape} {β : Type} (builder : Builder (Sequential σ τ)) (continuation : Sequential σ τ → IO β) : IO β

Build a model using the next global seed, then run a continuation.

nn.Sequential lives in Type 1, so executable code cannot receive the model directly from IO. nn.buildIO covers most uses without a continuation; this form remains for code that already uses continuation-passing style.

Example:

def builder : nn.Builder (nn.Sequential [2] [1]) :=
  nn.Sequential![nn.linear 2 8, nn.relu, nn.linear 8 1]

-- `nn.Sequential` lives in `Type 1`, so it cannot be returned from `IO`. Drawing the seed in `IO`
-- and handing the model to a continuation keeps model building pure.
def main : IO Unit :=
  nn.withModel builder fun model => nn.printSummary model