Layer Implementations

Explicit-seed layer constructors behind the public nn.* builders in NN.API.Seeded.

Only NN.API.Seeded imports this module. Every constructor here takes its initialization seeds as ordinary arguments; the public builders draw those seeds from nn.Builder and are the supported way to construct layers. The bodies stay exposed so that downstream proofs can unfold a built model through the public builders.

Leading dimensions

def TorchLean.nn.Impl.adaptLeadingShape (leading : Shape) {σ τ : Shape} (layer : Layer (σ.prependDim leading.size) (τ.prependDim leading.size)) : Layer (leading.concat σ) (leading.concat τ)

Reshape arbitrary leading dimensions into the single outer dimension expected by layer.

For an input of shape leading.concat σ, the layer receives shape [leading.size].concat σ; its output is then reshaped from [leading.size].concat τ to leading.concat τ. The adapter reuses the layer's parameters and buffer-update function.

def TorchLean.nn.Impl.batchedRecurrent (batchShape : Shape) {sequenceLength inputWidth hiddenWidth : ℕ} (core : Sequential [sequenceLength, inputWidth] [sequenceLength, hiddenWidth]) : Sequential ((batchShape.appendDim sequenceLength).appendDim inputWidth) ((batchShape.appendDim sequenceLength).appendDim hiddenWidth)

Share a single-sequence recurrent core over every index of batchShape.

Affine and recurrent layers

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

Linear layer on the last axis (prefix-shape preserving).

PyTorch analogue: torch.nn.Linear. See https://pytorch.org/docs/stable/generated/torch.nn.Linear.html.

If input has shape [..., inputWidth], linear inputWidth outputWidth returns a model of shape [..., outputWidth]. The leading dimensions are treated as a batch: they are flattened to (numel(prefix), inputWidth), the affine map is applied once, and the result is reshaped back.

def TorchLean.nn.Impl.rnn (sequenceLength inputWidth hiddenWidth : ℕ) (weightSeed : ℕ := 0) : Sequential [sequenceLength, inputWidth] [sequenceLength, hiddenWidth]

Vanilla RNN layer (time-major sequence, no batch axis).

Semantics:

$$ h_t=\tanh\!\left(W[x_t;h_{t-1}]+b\right),\qquad h_{-1}=0. $$

This is implemented by unrolling sequenceLength steps using existing TorchLean ops, so it runs on both CPU and CUDA backends.

PyTorch analogy: torch.nn.RNN(inputWidth, hiddenWidth, nonlinearity="tanh") with batch_first=false, specialized to a single batch element.

def TorchLean.nn.Impl.gru (sequenceLength inputWidth hiddenWidth : ℕ) (resetWeightSeed updateWeightSeed candidateWeightSeed : ℕ := 0) : Sequential [sequenceLength, inputWidth] [sequenceLength, hiddenWidth]

GRU layer (time-major sequence, no batch axis).

This is implemented by unrolling sequenceLength Cho-style steps using existing TorchLean ops, so it runs on both CPU and CUDA backends. PyTorch uses a different reset-after candidate parameterization; its GRU checkpoints are not directly compatible with this constructor.

def TorchLean.nn.Impl.mamba (sequenceLength inputWidth hiddenWidth : ℕ) (inputWeightSeed stateWeightSeed gateWeightSeed : ℕ := 0) (options : Runtime.Autograd.Model.Mamba.Options := { }) : Sequential [sequenceLength, inputWidth] [sequenceLength, hiddenWidth]

Trainable selective Mamba layer.

The input has shape (sequenceLength × inputWidth) and the output has shape (sequenceLength × hiddenWidth). Each token passes through a causal depthwise convolution and produces its own time steps, input coefficients, and readout coefficients for a diagonal state update. options controls the expanded channels, states per channel, and convolution width. The recurrence is unrolled with differentiable tensor operations.

def TorchLean.nn.Impl.lstm (sequenceLength inputWidth hiddenWidth : ℕ) (forgetWeightSeed inputWeightSeed candidateWeightSeed outputWeightSeed : ℕ := 0) : Sequential [sequenceLength, inputWidth] [sequenceLength, hiddenWidth]

LSTM layer (time-major sequence, no batch axis).

This is implemented by unrolling sequenceLength steps using existing TorchLean ops, so it runs on both CPU and CUDA backends.

PyTorch analogy: torch.nn.LSTM(inputWidth, hiddenWidth) with batch_first=false, specialized to a single batch element.

Shape and reduction layers

def TorchLean.nn.Impl.softmax {s : Shape} (axis : ℕ) : Sequential s s

Softmax over a tensor dimension, rejected by model validation when the axis is out of bounds.

def TorchLean.nn.Impl.logSoftmax {s : Shape} (axis : ℕ) : Sequential s s

Stable log-softmax over a tensor dimension, rejected by model validation when the axis is out of bounds.

def TorchLean.nn.Impl.sum {s : Shape} : Sequential s []

Reduce-sum to a scalar. PyTorch analogue: torch.sum.

def TorchLean.nn.Impl.flatten {s : Shape} : Sequential s [s.size]

Flatten any tensor into a 1D vector of length size s. PyTorch analogue: torch.flatten.

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

View a tensor with a new shape containing the same number of scalar entries.

This is the shape-typed counterpart of torch.reshape. A size mismatch is represented as an invalid model configuration and rejected by the ordinary validation path.

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

Flatten each tensor after an arbitrary batch shape.

For batchShape = [batch], this is the typed counterpart of torch.flatten(x, start_dim=1). Multiple batch dimensions are preserved without introducing a separate batched tensor type.

def TorchLean.nn.Impl.dropout {s : Shape} (p : Float) (seed : ℕ := 0) : Sequential s s

Dropout layer (active in train mode, identity in eval mode).

PyTorch analogue: torch.nn.Dropout.

Convolution and pooling

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

Apply an arbitrary-dimensional convolution to the channel and spatial suffix of a tensor.

The input suffix is (inputChannels, spatial...). Any axes in batchShape are preserved; internally they are flattened into one runtime batch and restored after the convolution.

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

Apply an arbitrary-dimensional transpose convolution to the channel and spatial suffix.

The input suffix is (inputChannels, spatial...). Any axes in batchShape are mapped independently and restored after the operation.

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

Apply max pooling to the channel and spatial suffix of a tensor.

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

Apply average pooling to the channel and spatial suffix of a tensor.

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

Global average pooling over every spatial axis, preserving the batch axes and channels.

Normalization

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

Layer normalization over the final axis of a tensor.

Every index in batchShape selects one vector of length width. Its entries share a mean and variance, while scale and bias are shared across all leading indices. batchShape := [] describes a single vector, and zero-sized leading axes are also allowed. Only width must be positive.

eps is added to the variance before the square root. Setting bias := false keeps only the learned scale; setting affine := false removes both learned parameters.

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

Divide each final-axis vector by sqrt(mean(x * x) + eps), then apply a learned scale.

The scale has shape [width] and is shared across all leading indices. Setting affine := false removes it from model state. The default eps is 1e-5, independent of the scalar type.

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

Batch normalization over (batchShape..., channels, spatial...) for any spatial rank.

All leading batch axes and spatial axes contribute to each channel's training statistics. Evaluation uses the running mean and variance. momentum controls their moving-average updates, and eps is added to the variance before taking the square root.

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

Instance normalization over (batchShape..., channels, spatial...) for any spatial rank.

Each sample and channel uses its own spatial mean and variance in both training and evaluation. Scale and bias are shared across samples. Setting bias := false keeps only the scale, while affine := false removes both parameters.

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

Group normalization over (batchShape..., channels, spatial...) for any spatial rank.

Within each sample, each group shares a mean and variance across its channels and spatial positions. Scale and bias have one entry per channel. Setting bias := false keeps only the scale; setting affine := false removes both parameters.

Attention

def TorchLean.nn.Impl.multiHeadAttention (batchShape : Shape := []) {sequenceLength modelWidth : ℕ} (config : MultiHeadAttention.Config) (queryWeightSeed keyWeightSeed valueWeightSeed outputWeightSeed : ℕ := 0) (mask : Option (Tensor Bool [sequenceLength, sequenceLength]) := none) (dropoutSeed : ℕ := 0) : Sequential (batchShape.concat [sequenceLength, modelWidth]) (batchShape.concat [sequenceLength, modelWidth])

Multi-head self-attention over a trailing (sequenceLength × modelWidth) shape.

If mask is provided, it is a boolean attention mask of shape (n × n) (e.g. causal masking).

Positional encodings

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

Add learned positional embeddings to the (sequenceLength × embeddingWidth) suffix of a tensor.

PyTorch analogue: x + position[:sequenceLength] where position is a parameter table.

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

Add sinusoidal positional encodings to the (sequenceLength × embeddingWidth) suffix of a tensor.

Implementation:

• precompute PE : (sequenceLength × embeddingWidth) at initialization time (stored as a non-trainable buffer),
• broadcast it across batchShape and add it to the input.

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

Apply RoPE to the (sequenceLength × headWidth) suffix of a tensor.

This matches the standard identity:

$$ \operatorname{rope}(x) = x \odot \cos + \operatorname{rotatePairs}(x) \odot \sin $$

where cos and sin depend only on (pos, dim) and broadcast across batchShape.

Notes:

• This layer is differentiable (gradients flow through the rotation), but it has no trainable parameters; the precomputed cos/sin tables are stored as non-trainable buffers.
• The pure spec version is in NN.Spec.Layers.PositionalEncoding (Spec.ropeApplyHeadsSpec).

Embeddings

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

Linear projection for one-hot or soft token-distribution inputs.

Input shape: [..., vocabularySize] Output shape: [..., embeddingWidth]

This is not an indexed embedding: it multiplies the final input axis by a trainable table. Use embedding for bounded token ids.

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

Build a trainable table for bounded token ids.

Blocks and heads

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

Build a rank-polymorphic convolution/activation block.

def TorchLean.nn.Impl.transformerEncoderBlock (batchShape : Shape := []) {sequenceLength modelWidth : ℕ} (config : TransformerEncoder.Block.Config) (queryWeightSeed keyWeightSeed valueWeightSeed outputWeightSeed firstFeedForwardWeightSeed secondFeedForwardWeightSeed attentionDropoutSeed feedForwardDropoutSeed : ℕ := 0) (mask : Option (Tensor Bool [sequenceLength, sequenceLength]) := none) (attentionProbabilityDropoutSeed feedForwardHiddenDropoutSeed : ℕ := 0) : Sequential (batchShape.concat [sequenceLength, modelWidth]) (batchShape.concat [sequenceLength, modelWidth])

Transformer encoder block.

With post-normalization (the default), this follows: LayerNorm(x + MHA(x)) -> LayerNorm(x + FFN(x)).

With normalizeFirst := true, each branch is normalized before its learned transform: x + MHA(LayerNorm(x)) -> x + FFN(LayerNorm(x)).

PyTorch analogue:

• torch.nn.TransformerEncoderLayer (https://pytorch.org/docs/stable/generated/torch.nn.TransformerEncoderLayer.html)

def TorchLean.nn.Impl.affineHead (batchShape : Shape := []) {featureShape : Shape} (outputWidth : ℕ) (weightSeed : ℕ := 0) : Sequential (batchShape.concat featureShape) (batchShape.appendDim outputWidth)

Flatten the feature suffix and apply an affine map, preserving every batch axis.