Transformer Blocks

This module exposes attention, feed-forward, and Transformer-stack constructors used by sequence models and higher-level examples.

structure TorchLean.nn.Internal.blocks.TransformerEncoderBlock : Type

Config record for transformerEncoderBlock.

Separating the config as a structure makes it easier to write readable examples and keep seed management deterministic.

• numHeads : ℕ

  Number of attention heads.

• headDim : ℕ

  Per-head embedding dimension.

• ffnHidden : ℕ

  Hidden dimension of the feed-forward network.

• activation : Activation

  Activation used in the feed-forward network.

• dropout? : Option Float

  Optional dropout probability for examples; none means no dropout.

• normFirst : Bool

  Normalize before attention and feed-forward sublayers instead of after each residual.

• attentionOutputBias : Bool

  Add a trainable bias after the attention output projection.

• weightInit? : Option Runtime.Autograd.Torch.Init.Scheme

  Attention and feed-forward weight initialization. none keeps each layer's default.

• residualOutputInit? : Option Runtime.Autograd.Torch.Init.Scheme

  Initializer for the attention and feed-forward projections that write to residual streams.

  When omitted, weightInit? is used. The separate field supports depth-scaled residual initialization without imposing that convention on every Transformer.

• seedBase : ℕ

  Base seed used to derive deterministic per-layer seeds inside the block.

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

Transformer encoder block configuration.

This follows the familiar pattern: (residual MHA) -> LayerNorm -> (residual FFN) -> LayerNorm.

PyTorch analogue:

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

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

Transformer encoder block.

This is transformerEncoderBlockWithMask; pass mask := ... to enable causal masking (or other attention masks).

structure TorchLean.nn.Internal.blocks.TransformerEncoderStack : Type

Config record for transformerEncoderStack.

This builds layers copies of transformerEncoderBlock, allocating seeds in a fixed stride.

• layers : ℕ

  Layer stack.

• block : TransformerEncoderBlock

  Template config for each block (its seedBase is ignored; we allocate per-layer seeds).

• seedBase : ℕ

  Base seed for the whole stack.

• seedStride : ℕ

  Seed stride between consecutive blocks (must exceed the per-block seed footprint).

def TorchLean.nn.Internal.blocks.transformerStackGoWithMask {batch n dModel : ℕ} [NeZero n] [NeZero dModel] (template : TransformerEncoderBlock) (seedBase seedStride : ℕ) (mask : Option (Spec.Tensor Bool (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) := none) (layerIdx remaining : ℕ) : 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)))

Internal recursion for transformerEncoderStack.

Builds remaining blocks starting at layerIdx, allocating each block's seedBase as seedBase + layerIdx * seedStride.

def TorchLean.nn.Internal.blocks.transformerStackGo {batch n dModel : ℕ} [NeZero n] [NeZero dModel] (template : TransformerEncoderBlock) (seedBase seedStride layerIdx remaining : ℕ) : 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)))

Internal recursion for transformerEncoderStack (unmasked).

This is transformerStackGoWithMask with mask := none.

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

Stack cfg.layers copies of blocks.transformerEncoderBlock.

TorchLean analogue of composing torch.nn.TransformerEncoderLayer into a torch.nn.TransformerEncoder, using Seq composition for the typed model.

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

Stack cfg.layers copies of blocks.transformerEncoderBlock.

This is transformerEncoderStackWithMask; pass mask := ... to enable causal masking (or other attention masks).

def TorchLean.nn.Internal.blocks.transformerEncoderClassifier {batch n dModel : ℕ} [NeZero n] [NeZero dModel] (classes : ℕ) (cfg : TransformerEncoderStack) : Sequential (Spec.Shape.dim batch (Spec.Shape.dim n (Spec.Shape.dim dModel Spec.Shape.scalar))) (Spec.Shape.dim batch (Spec.Shape.dim classes Spec.Shape.scalar))

Transformer encoder followed by a flatten+linear classification head.

PyTorch analogue (approximately): nn.TransformerEncoder(...) + pooling/flattening + nn.linear.