Seeded model builders

Layer constructors draw deterministic initialization seeds from an explicit seed stream.

Model Builders and Seeding

TorchLean keeps initialization randomness explicit so examples are reproducible.

Layer constructors return nn.M, a deterministic state computation over the initialization seed stream. Call nn.run seed to construct a model reproducibly.

Note: nn.Sequential lives in Type 2, so it cannot be returned directly from IO. We keep model building pure by drawing a base seed in IO and then calling nn.run.

def TorchLean.nn.manualSeed (seed : ℕ) : IO Unit

Set the global seed used by nn.runGlobal and nn.nextSeed.

Prefer nn.run seed when the seed belongs in the model definition itself.

def TorchLean.nn.globalAvgPool (leading : Spec.Shape := Spec.Shape.scalar) {d channels : ℕ} (spatial : Vector ℕ d) (spatialNonzero : ∀ (i : Fin d), spatial.get i ≠ 0) : rand.SeedM (Sequential (leading.concat (Spec.Shape.ofList (channels :: spatial.toList))) (leading.concat (Spec.Shape.dim channels Spec.Shape.scalar)))

Build global average pooling over the supplied nonempty spatial dimensions without consuming a seed.

Default Builders

The nn.* constructors allocate initialization seeds through nn.M.

@[reducible, inline]

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

Seeded builder monad: a state monad over TorchLean.rand.SeedStream.

def TorchLean.nn.run {α : Type 2} (seed : ℕ) (x : M α) : α

Run a seeded builder starting from a base seed.

def TorchLean.nn.lift {α : Type 2} (x : α) : M α

Lift a pure value into the seeded builder (consumes no seeds).

def TorchLean.nn.mapLeading (leading : Spec.Shape) {σ τ : Spec.Shape} (model : Sequential σ τ) : M (Sequential (leading.concat σ) (leading.concat τ))

Apply a model independently over an arbitrary collection of leading dimensions.

def TorchLean.nn.withSeed {α : Type 2} (k : ℕ → α) : M α

Consume one fresh seed and pass it to k.

def TorchLean.nn.withSeedPair {α : Type 2} (k : ℕ → ℕ → α) : M α

Consume two fresh seeds and pass them to k (in order).

def TorchLean.nn.relu {s : Spec.Shape} : M (Sequential s s)

Build an elementwise ReLU layer without consuming an initialization seed.

def TorchLean.nn.silu {s : Spec.Shape} : M (Sequential s s)

Build an elementwise SiLU layer without consuming an initialization seed.

def TorchLean.nn.gelu {s : Spec.Shape} : M (Sequential s s)

Build an elementwise GELU layer without consuming an initialization seed.

def TorchLean.nn.sigmoid {s : Spec.Shape} : M (Sequential s s)

Build an elementwise sigmoid layer without consuming an initialization seed.

def TorchLean.nn.tanh {s : Spec.Shape} : M (Sequential s s)

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

def TorchLean.nn.softmax {s : Spec.Shape} : M (Sequential s s)

Build a softmax layer over a tensor shape without consuming an initialization seed.

def TorchLean.nn.sum {s : Spec.Shape} : M (Sequential s Spec.Shape.scalar)

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

def TorchLean.nn.flatten {s : Spec.Shape} : M (Sequential s (Spec.Shape.dim s.size Spec.Shape.scalar))

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

def TorchLean.nn.flattenBatch {n : ℕ} {s : Spec.Shape} : M (Sequential (Spec.Shape.dim n s) (Spec.Shape.dim n (Spec.Shape.dim s.size Spec.Shape.scalar)))

Build a layer that preserves the batch axis while flattening each example.

def TorchLean.nn.maxPool (leading : Spec.Shape := Spec.Shape.scalar) {d channels : ℕ} (spatial : Vector ℕ d) (cfg : Pool d) : M (Sequential (leading.concat (Spec.Shape.ofList (channels :: spatial.toList))) (leading.concat (Spec.Shape.ofList (channels :: (Spec.poolOutSpatialPad spatial cfg.kernel cfg.stride cfg.padding).toList))))

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

def TorchLean.nn.avgPool (leading : Spec.Shape := Spec.Shape.scalar) {d channels : ℕ} (spatial : Vector ℕ d) (cfg : Pool d) : M (Sequential (leading.concat (Spec.Shape.ofList (channels :: spatial.toList))) (leading.concat (Spec.Shape.ofList (channels :: (Spec.poolOutSpatialPad spatial cfg.kernel cfg.stride cfg.padding).toList))))

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

def TorchLean.nn.linear (inDim outDim : ℕ) (pfx : Spec.Shape := Spec.Shape.scalar) : M (Sequential (pfx.appendDim inDim) (pfx.appendDim outDim))

Build an affine layer, consuming independent seeds for its weight and bias initializers.

def TorchLean.nn.linearWith (inDim outDim : ℕ) (cfg : Linear) (pfx : Spec.Shape := Spec.Shape.scalar) : M (Sequential (pfx.appendDim inDim) (pfx.appendDim outDim))

Seeded affine layer with an explicit initialization policy.

def TorchLean.nn.linearV (inDim outDim : ℕ) : M (Sequential (Spec.Shape.dim inDim Spec.Shape.scalar) (Spec.Shape.dim outDim Spec.Shape.scalar))

Vector-only linear layer, specialized to the scalar prefix shape.

def TorchLean.nn.deterministic.linear (inDim outDim seedWeight seedBias : ℕ) (leading : Spec.Shape := Spec.Shape.scalar) : Sequential (leading.appendDim inDim) (leading.appendDim outDim)

Construct a linear layer with explicit parameter-initialization seeds.

def TorchLean.nn.rnn (seqLen inputSize hiddenSize : ℕ) : M (Sequential (Spec.Shape.dim seqLen (Spec.Shape.dim inputSize Spec.Shape.scalar)) (Spec.Shape.dim seqLen (Spec.Shape.dim hiddenSize Spec.Shape.scalar)))

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

def TorchLean.nn.gru (seqLen inputSize hiddenSize : ℕ) : M (Sequential (Spec.Shape.dim seqLen (Spec.Shape.dim inputSize Spec.Shape.scalar)) (Spec.Shape.dim seqLen (Spec.Shape.dim hiddenSize Spec.Shape.scalar)))

Build a seeded gated recurrent unit over a fixed sequence length.

def TorchLean.nn.mamba (seqLen inputSize hiddenSize : ℕ) : M (Sequential (Spec.Shape.dim seqLen (Spec.Shape.dim inputSize Spec.Shape.scalar)) (Spec.Shape.dim seqLen (Spec.Shape.dim hiddenSize Spec.Shape.scalar)))

Build a seeded Mamba-style state-space sequence layer.

def TorchLean.nn.lstm (seqLen inputSize hiddenSize : ℕ) : M (Sequential (Spec.Shape.dim seqLen (Spec.Shape.dim inputSize Spec.Shape.scalar)) (Spec.Shape.dim seqLen (Spec.Shape.dim hiddenSize Spec.Shape.scalar)))

Build a seeded long short-term memory layer over a fixed sequence length.

def TorchLean.nn.conv (leading : Spec.Shape := Spec.Shape.scalar) {d inChannels : ℕ} (spatial : Vector ℕ d) (cfg : Conv d) [NeZero inChannels] : M (Sequential (leading.concat (Spec.Shape.ofList (inChannels :: spatial.toList))) (leading.concat (Spec.Shape.ofList (cfg.outChannels :: (Spec.convOutSpatial spatial cfg.kernel cfg.stride cfg.padding).toList))))

Build an arbitrary-rank convolution, allocating separate kernel and bias seeds.

def TorchLean.nn.batchNorm (leading : Spec.Shape := Spec.Shape.scalar) {d channels : ℕ} (spatial : Vector ℕ d) (cfg : ChannelNorm := { }) [NeZero leading.size] [NeZero channels] [NeZero (Spec.Shape.ofList spatial.toList).size] : M (Sequential (leading.concat (Spec.Shape.ofList (channels :: spatial.toList))) (leading.concat (Spec.Shape.ofList (channels :: spatial.toList))))

Build batch normalization with seeded scale and offset parameters.

def TorchLean.nn.instanceNorm (leading : Spec.Shape := Spec.Shape.scalar) {d channels : ℕ} (spatial : Vector ℕ d) (cfg : ChannelNorm := { }) [NeZero leading.size] [NeZero channels] [NeZero (Spec.Shape.ofList spatial.toList).size] : M (Sequential (leading.concat (Spec.Shape.ofList (channels :: spatial.toList))) (leading.concat (Spec.Shape.ofList (channels :: spatial.toList))))

Build instance normalization with seeded scale and offset parameters.

def TorchLean.nn.groupNorm (leading : Spec.Shape := Spec.Shape.scalar) {d channels : ℕ} (spatial : Vector ℕ d) (groups : ℕ) (hGroups : groups > 0) (hGroupsLe : channels ≥ groups) (hDiv : channels % groups = 0) (cfg : ChannelNorm := { }) [NeZero leading.size] [NeZero channels] [NeZero (Spec.Shape.ofList spatial.toList).size] : M (Sequential (leading.concat (Spec.Shape.ofList (channels :: spatial.toList))) (leading.concat (Spec.Shape.ofList (channels :: spatial.toList))))

Build group normalization after checking the positive group count and channel divisibility.

def TorchLean.nn.embedding (vocab embedDim : ℕ) (cfg : Embedding := { }) {pfx : Spec.Shape} : M (Sequential (pfx.appendDim vocab) (pfx.appendDim embedDim))

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

def TorchLean.nn.sinusoidalPositionalEncoding {batch seqLen embedDim : ℕ} (cfg : SinusoidalPositionalEncoding := { }) : M (Sequential (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))))

Build deterministic sinusoidal positional encoding for a batched sequence.

def TorchLean.nn.rope {batch numHeads seqLen headDim : ℕ} (cfg : RoPE := { }) : M (Sequential (Spec.Shape.dim batch (Spec.Shape.dim numHeads (Spec.Shape.dim seqLen (Spec.Shape.dim headDim Spec.Shape.scalar)))) (Spec.Shape.dim batch (Spec.Shape.dim numHeads (Spec.Shape.dim seqLen (Spec.Shape.dim headDim Spec.Shape.scalar)))))

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

def TorchLean.nn.learnedPositionalEmbedding {batch seqLen embedDim : ℕ} (cfg : LearnedPositionalEmbedding := { }) : M (Sequential (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))))

Build learned positional embeddings from a freshly allocated parameter seed.

def TorchLean.nn.layerNorm {batch seqLen embedDim : ℕ} [NeZero seqLen] [NeZero embedDim] : M (Sequential (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))))

Build layer normalization with independently seeded scale and offset parameters.

def TorchLean.nn.multiheadAttention {batch n dModel : ℕ} [NeZero n] (cfg : MultiheadAttention) (mask : Option (Spec.Tensor Bool (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) := none) : M (Sequential (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))))

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

def TorchLean.nn.transformerEncoderBlock {batch n dModel : ℕ} [NeZero n] [NeZero dModel] (cfg : blocks.TransformerEncoderBlock) (mask : Option (Spec.Tensor Bool (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) := none) : M (Sequential (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))))

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

def TorchLean.nn.transformerEncoderStack {batch n dModel : ℕ} [NeZero n] [NeZero dModel] (cfg : blocks.TransformerEncoderStack) (mask : Option (Spec.Tensor Bool (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) := none) : M (Sequential (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))))

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

def TorchLean.nn.dropout {s : Spec.Shape} (p : Float) : M (Sequential s s)

Build dropout with a fresh deterministic mask seed from the builder stream.

def TorchLean.nn.runGlobal {α : Type} (x : M α) : IO α

Run a seeded builder using the global seed stream set by nn.manualSeed (results in Type).

Note: model values like nn.Sequential live in Type 2, so they cannot be returned from IO. For models, use nn.run with an explicit base seed (obtained from nn.nextSeed).

def TorchLean.nn.nextSeed : IO ℕ

Draw a fresh base seed from the global seed stream set by nn.manualSeed.

def TorchLean.nn.nextSeeds (n : ℕ) : IO (List ℕ)

Draw n fresh base seeds from the global seed stream.

Naming Convenience

nn.run / nn.nextSeed are the core primitives, but in user code it is often clearer to read:

• “build a model from this seed” (nn.run)
• “draw a fresh init seed” (nn.freshSeed)
• “build a model using the next global init seed” (nn.withModel)

def TorchLean.nn.freshSeed : IO ℕ

Draw a fresh base seed from the global seed stream.

def TorchLean.nn.withModel {σ τ : Spec.Shape} {β : Type} (mk : M (Sequential σ τ)) (k : Sequential σ τ → IO β) : IO β

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

nn.Sequential lives in Type 2, so executable code passes the model to a continuation rather than returning it directly from IO.