Gated Recurrent Models

TorchLean provides GRU layers/cells in NN.Spec.Layers.Gru. This file builds models on top of that layer API: common compositions, heads, and a couple of end-to-end forward/backward routines.

Higher‑level GRU architectures built from module specs (Spec.Module.Chain):

• sequence‑to‑sequence outputs,
• classifier heads (many‑to‑one),
• multi‑layer compositions.

GRU cell equations are in NN/Spec/Layers/Gru.lean; this file is primarily “wiring”.

References:

• Cho et al. (2014), "Learning Phrase Representations using RNN Encoder–Decoder for Statistical Machine Translation" (introduces GRU): https://arxiv.org/abs/1406.1078
• Chung et al. (2014), "Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling" (GRU variants/ablation): https://arxiv.org/abs/1412.3555
• PyTorch nn.GRUCell docs: https://docs.pytorch.org/docs/stable/generated/torch.nn.GRUCell.html
• PyTorch nn.GRU docs: https://pytorch.org/docs/stable/generated/torch.nn.GRU.html

The wiring matches torch.nn.GRU with linear heads and pooling over time (e.g. last hidden state for classification), and each definition below names its wiring analogue. The comparison is architectural only. The recurrent core is the original Cho reset-before GRU (GRUSpec), while PyTorch uses the reset-after candidate equation (GRUResetAfterSpec). For general recurrent matrices these are different function families, so PyTorch GRU weights cannot be converted into these models.

def Spec.Gru.sequence {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {seqLen inputSize hiddenSize outputSize : ℕ} (gruSpec : GRUSpec α inputSize hiddenSize) (linearSpec : LinearSpec α hiddenSize outputSize) : Module.Chain α [seqLen, inputSize] [seqLen, outputSize]

A sequence-to-sequence GRU model, written as a Spec.Module.Chain.

Pipeline: GRU(seqLen, inputSize → hiddenSize) then Linear applied at each timestep.

Wiring analogue: nn.GRU(..., batch_first=False) followed by an nn.linear on the output sequence.

def Spec.Gru.classifier {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {seqLen inputSize hiddenSize numClasses : ℕ} (gruSpec : GRUSpec α inputSize hiddenSize) (classifierHead : LinearSpec α hiddenSize numClasses) (h : seqLen ≠ 0) : Module.Chain α [seqLen, inputSize] [numClasses]

A many-to-one GRU classifier (use the last hidden state, then a linear head).

Wiring analogue: run nn.GRU over the sequence and feed the last output/hidden state into nn.Linear(hiddenSize, numClasses).

def Spec.Gru.stacked {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {seqLen inputSize hiddenSize outputSize : ℕ} (layers : RecurrentStack (GRUSpec α) inputSize hiddenSize) (linearSpec : LinearSpec α hiddenSize outputSize) : Module.Chain α [seqLen, inputSize] [seqLen, outputSize]

A recurrent stack of arbitrary depth and widths, followed by a per-timestep linear head.

def Spec.Gru.languageModel {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {seqLen vocabularySize hiddenSize : ℕ} (embeddingSpec : LinearSpec α vocabularySize hiddenSize) (gruSpec : GRUSpec α hiddenSize hiddenSize) (outputSpec : LinearSpec α hiddenSize vocabularySize) : Module.Chain α [seqLen, vocabularySize] [seqLen, vocabularySize]

A simple GRU language-model style pipeline:

Linear as the embedding/projection map, then GRU, then a per-timestep projection back to vocabularySize.

Wiring analogue: embedding (often nn.Embedding), nn.GRU, and nn.Linear(hiddenSize, vocabularySize). We use LinearSpec here as a spec-friendly stand-in for a one-hot embedding matrix.

Record-style model specs

The Spec.Module.Chain builders above are the most uniform way to assemble models in TorchLean.

The declarations below use small record types with explicit forward functions. This is useful when you want to talk about a particular architecture directly (e.g. encoder-decoder), or when you need to carry extra per-model parameters (e.g. a dropout rate) without building a full module stack.

structure Spec.Gru.Model (α : Type) [TorchLean.Storage α] (inputSize hiddenSize outputSize : ℕ) : Type

Bundle of parameters for a single-layer GRU model with a linear output head.

This is a direct record representation (as opposed to the Spec.Module.Chain representation above).

• gru : GRUSpec α inputSize hiddenSize

  Recurrent cell parameters.

• outputLayer : LinearSpec α hiddenSize outputSize

  Linear output projection.

structure Spec.Gru.CellGrads (α : Type) [TorchLean.Storage α] (inputSize hiddenSize : ℕ) : Type

Gradients of the three GRU gates.

• resetWeight : TorchLean.Tensor α [hiddenSize, inputSize + hiddenSize]

  Gradient of the reset-gate weight matrix.

• resetBias : TorchLean.Tensor α [hiddenSize]

  Gradient of the reset-gate bias.

• updateWeight : TorchLean.Tensor α [hiddenSize, inputSize + hiddenSize]

  Gradient of the update-gate weight matrix.

• updateBias : TorchLean.Tensor α [hiddenSize]

  Gradient of the update-gate bias.

• candidateWeight : TorchLean.Tensor α [hiddenSize, inputSize + hiddenSize]

  Gradient of the candidate-state weight matrix.

• candidateBias : TorchLean.Tensor α [hiddenSize]

  Gradient of the candidate-state bias.

structure Spec.Gru.Grads (α : Type) [TorchLean.Storage α] (inputSize hiddenSize outputSize : ℕ) : Type

Parameter gradients for a GRU model and its linear output head.

• cell : CellGrads α inputSize hiddenSize

  Gradients of the recurrent cell.

• outputWeight : TorchLean.Tensor α [outputSize, hiddenSize]

  Gradient of the output projection weight.

• outputBias : TorchLean.Tensor α [outputSize]

  Gradient of the output projection bias.

structure Spec.Gru.StackedModel (α : Type) [TorchLean.Storage α] (inputSize hiddenSize outputSize : ℕ) : Type

Recurrent cells with independently chosen widths and a linear output head.

The endpoint hiddenSize is the width of the last cell. An empty stack has hiddenSize = inputSize.

• layers : RecurrentStack (GRUSpec α) inputSize hiddenSize

  Layer dimensions determine the type of every intermediate hidden stream and state.

• outputLayer : LinearSpec α hiddenSize outputSize

  Linear output projection.

def Spec.Gru.StackedModel.forward {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (model : StackedModel α inputSize hiddenSize outputSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (initialHiddens : model.layers.States fun (width : ℕ) => TorchLean.Tensor α [width]) : TorchLean.Tensor α [seqLen, outputSize] × model.layers.States fun (width : ℕ) => TorchLean.Tensor α [width]

Run every layer with its own initial state, then project the final hidden stream.

The result retains all final states. Empty sequences leave each initial state unchanged, and an empty stack applies the head directly to the input sequence.

structure Spec.Gru.Classifier (α : Type) [TorchLean.Storage α] (inputSize hiddenSize numClasses : ℕ) : Type

Bundle of parameters for a many-to-one GRU classifier.

The classifier head is applied to the final hidden state.

• gru : GRUSpec α inputSize hiddenSize

  Recurrent cell parameters.

• classifier : LinearSpec α hiddenSize numClasses

  Linear classifier head.

structure Spec.Gru.Generator (α : Type) [TorchLean.Storage α] (vocabularySize hiddenSize : ℕ) : Type

Bundle of parameters for a many-to-many GRU generator (language-model style).

This includes an (embedding) linear map, recurrent core, and output projection back to vocabulary.

• embedding : LinearSpec α vocabularySize hiddenSize

  Token projection used by this one-hot specification.

• gru : GRUSpec α hiddenSize hiddenSize

  Recurrent cell parameters.

• outputProjection : LinearSpec α hiddenSize vocabularySize

  Projection from hidden states to vocabulary logits.

structure Spec.Gru.BidirectionalModel (α : Type) [TorchLean.Storage α] (inputSize hiddenSize outputSize : ℕ) : Type

Bundle of parameters for a bidirectional GRU model with an output head.

The head consumes the concatenation of forward and backward hidden states. Wiring analogue: nn.GRU(..., bidirectional=true) plus a linear projection.

• forwardGru : GRUSpec α inputSize hiddenSize

  Recurrent cell for the original sequence order.

• backwardGru : GRUSpec α inputSize hiddenSize

  Recurrent cell for the reversed sequence order.

• outputLayer : LinearSpec α (hiddenSize + hiddenSize) outputSize

  Projection from concatenated forward and backward states.

structure Spec.Gru.LanguageModel (α : Type) [TorchLean.Storage α] (vocabularySize hiddenSize : ℕ) : Type

Bundle of parameters for a stacked GRU language model with deterministic dropout.

This model uses an array of GRU layers (all with hiddenSize input/output) and applies evaluation-mode dropout between the GRU stack and the output projection.

• embedding : LinearSpec α vocabularySize hiddenSize

  Token projection used by this one-hot specification.

• layers : Array (GRUSpec α hiddenSize hiddenSize)

  Recurrent layers, ordered from input to output.

• outputProjection : LinearSpec α hiddenSize vocabularySize

  Projection from hidden states to vocabulary logits.

• dropoutRate : α

  Dropout probability used between the recurrent stack and output projection.

structure Spec.Gru.EncoderDecoder (α : Type) [TorchLean.Storage α] (inputVocabSize hiddenSize outputVocabSize : ℕ) : Type

Bundle of parameters for a GRU encoder-decoder model (seq2seq).

This uses separate embeddings and GRU cores for encoder and decoder, plus an output projection. Wiring analogue: an encoder nn.GRU and a decoder nn.GRU with teacher forcing.

• encoderEmbedding : LinearSpec α inputVocabSize hiddenSize

  Source-token projection.

• encoderGru : GRUSpec α hiddenSize hiddenSize

  Encoder recurrent cell.

• decoderEmbedding : LinearSpec α outputVocabSize hiddenSize

  Target-token projection.

• decoderGru : GRUSpec α hiddenSize hiddenSize

  Decoder recurrent cell.

• outputProjection : LinearSpec α hiddenSize outputVocabSize

  Projection from decoder states to target-vocabulary logits.

def Spec.Gru.Model.forward {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize outputSize : ℕ} (model : Model α inputSize hiddenSize outputSize) (input : TorchLean.Tensor α [inputSize]) (hidden : TorchLean.Tensor α [hiddenSize]) : TorchLean.Tensor α [outputSize] × TorchLean.Tensor α [hiddenSize]

One-step forward for Gru.Model.

Input: (x_t, h_{t-1}). Output: (y_t, h_t).

def Spec.Gru.Model.forwardSequence {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (model : Model α inputSize hiddenSize outputSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (initialHidden : TorchLean.Tensor α [hiddenSize]) : TorchLean.Tensor α [seqLen, outputSize] × TorchLean.Tensor α [hiddenSize]

Sequence forward for Gru.Model (time-major).

Returns (outputs, final_hidden). Empty sequences preserve the initial hidden state.

Wiring analogue: run nn.GRU over the sequence, then apply nn.linear at each timestep.

def Spec.Gru.Classifier.forward {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize numClasses : ℕ} (model : Classifier α inputSize hiddenSize numClasses) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (initialHidden : TorchLean.Tensor α [hiddenSize]) (h : 0 < seqLen) : TorchLean.Tensor α [numClasses]

Forward pass for a Gru.Classifier (many-to-one).

This runs the GRU over the input sequence and applies the classifier head to the final hidden state.

def Spec.Gru.Generator.forward {α : Type} [TorchLean.Storage α] [Context α] {seqLen vocabularySize hiddenSize : ℕ} (model : Generator α vocabularySize hiddenSize) (inputTokens : TorchLean.Tensor α [seqLen, vocabularySize]) (initialHidden : TorchLean.Tensor α [hiddenSize]) : TorchLean.Tensor α [seqLen, vocabularySize] × TorchLean.Tensor α [hiddenSize]

Forward pass for a Gru.Generator (many-to-many).

This applies an embedding linear map to each token vector, runs the GRU, and projects each hidden state back into vocabulary space. Empty sequences preserve the initial hidden state.

def Spec.Gru.BidirectionalModel.forward {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (model : BidirectionalModel α inputSize hiddenSize outputSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (forwardHidden backwardHidden : TorchLean.Tensor α [hiddenSize]) : TorchLean.Tensor α [seqLen, outputSize]

Forward pass for a bidirectional GRU model (time-major).

This runs a forward GRU on the sequence, a backward GRU on the reversed sequence, concatenates the two hidden streams per timestep, and applies an output head.

def Spec.Gru.LanguageModel.forward {α : Type} [TorchLean.Storage α] [Context α] {seqLen vocabularySize hiddenSize : ℕ} (model : LanguageModel α vocabularySize hiddenSize) (inputTokens : TorchLean.Tensor α [seqLen, vocabularySize]) (initialHiddens : Array (TorchLean.Tensor α [hiddenSize])) (h : 0 < seqLen) : Option (TorchLean.Tensor α [seqLen, vocabularySize] × Array (TorchLean.Tensor α [hiddenSize]))

Forward pass for Gru.LanguageModel (teacher forcing, time-major).

This runs the embedding, then a stack of GRU layers with provided initial hiddens, applies evaluation-mode dropout (dropoutInferenceSpec), and projects to vocabulary logits.

def Spec.Gru.LanguageModel.forward.processLayers {α : Type} [TorchLean.Storage α] [Context α] {seqLen hiddenSize : ℕ} (initialHiddens : Array (TorchLean.Tensor α [hiddenSize])) (h : 0 < seqLen) (layers : List (GRUSpec α hiddenSize hiddenSize)) (index : ℕ) (layerInput : TorchLean.Tensor α [seqLen, hiddenSize]) : Option (TorchLean.Tensor α [seqLen, hiddenSize] × Array (TorchLean.Tensor α [hiddenSize]))

def Spec.Gru.EncoderDecoder.forward {α : Type} [TorchLean.Storage α] [Context α] {srcSeqLen tgtSeqLen inputVocabSize hiddenSize outputVocabSize : ℕ} (model : EncoderDecoder α inputVocabSize hiddenSize outputVocabSize) (sourceTokens : TorchLean.Tensor α [srcSeqLen, inputVocabSize]) (targetTokens : TorchLean.Tensor α [tgtSeqLen, outputVocabSize]) (encoderHidden : TorchLean.Tensor α [hiddenSize]) (hSource : 0 < srcSeqLen) (hTarget : 0 < tgtSeqLen) : TorchLean.Tensor α [tgtSeqLen, outputVocabSize] × TorchLean.Tensor α [hiddenSize] × TorchLean.Tensor α [hiddenSize]

Encoder-decoder forward pass (GRU encoder + GRU decoder).

This is a small reference architecture:

• encode sourceTokens into a final hidden state,
• decode targetTokens starting from that hidden state (teacher forcing),
• project decoder states into output-vocabulary logits.

Wiring analogue: nn.GRU encoder + nn.GRU decoder with a linear output projection.

def Spec.Gru.Model.backward {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (model : Model α inputSize hiddenSize outputSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (hiddenStates : TorchLean.Tensor α [seqLen, hiddenSize]) (outputGrad : TorchLean.Tensor α [seqLen, outputSize]) (resetGates updateGates candidates : TorchLean.Tensor α [seqLen, hiddenSize]) (h : seqLen ≠ 0) (initialHidden : TorchLean.Tensor α [hiddenSize] := TorchLean.Tensor.full [hiddenSize] 0) : Grads α inputSize hiddenSize outputSize × TorchLean.Tensor α [seqLen, inputSize]

Backward pass for Gru.Model using full backpropagation through time.

This assumes you already ran a forward pass that saved:

• hiddenStates,
• the GRU intermediates resetGates, updateGates, and candidates.

Those intermediates can be produced using Spec.gruExtractIntermediateValues from NN.Spec.Layers.Gru. Supply the same initialHidden as the forward pass; omitting it retains the zero-initial-state convention.

structure Spec.Gru.ResidualModel (α : Type) [TorchLean.Storage α] (inputSize hiddenSize outputSize : ℕ) : Type

Bundle of parameters for a residual GRU model.

This includes a projection from input space to hidden space so the input can be added as a residual to the GRU hidden stream.

• gru : GRUSpec α inputSize hiddenSize

  Recurrent cell parameters.

• residualProjection : LinearSpec α inputSize hiddenSize

  Projection that gives the input stream the recurrent hidden width.

• outputLayer : LinearSpec α hiddenSize outputSize

  Linear output projection.

def Spec.Gru.ResidualModel.forward {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (model : ResidualModel α inputSize hiddenSize outputSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (initialHidden : TorchLean.Tensor α [hiddenSize]) (h : 0 < seqLen) : TorchLean.Tensor α [seqLen, outputSize] × TorchLean.Tensor α [hiddenSize]

Forward pass for Gru.ResidualModel.

This runs the GRU, adds a projected version of the input as a residual connection, and applies the output head per timestep. The second result is the last residual feature hiddenStates[last] + projectedInputs[last]. For a continuation chunk, retain the raw GRU hidden state before this residual addition and use that as initialHidden.

def Spec.Gru.Model.toModule {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (model : Model α inputSize hiddenSize outputSize) : Module α [seqLen, inputSize] [seqLen, outputSize]

Package Gru.Model as a shape-indexed module.

The Python expression records the intended runtime analogue; forward remains the mathematical meaning of the module.

def Spec.Gru.Classifier.toModule {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize numClasses : ℕ} (model : Classifier α inputSize hiddenSize numClasses) (h : 0 < seqLen) : Module α [seqLen, inputSize] [numClasses]

Package Gru.Classifier as an Spec.Module.

Wiring analogue: nn.GRU feeding a nn.linear classifier head.

def Spec.Gru.BidirectionalModel.toModule {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (model : BidirectionalModel α inputSize hiddenSize outputSize) : Module α [seqLen, inputSize] [seqLen, outputSize]

Package Gru.BidirectionalModel as an Spec.Module.

Wiring analogue: nn.GRU(..., bidirectional=true) feeding a per-timestep linear head.

def Spec.Gru.Generator.toModule {α : Type} [TorchLean.Storage α] [Context α] {seqLen vocabularySize hiddenSize : ℕ} (model : Generator α vocabularySize hiddenSize) : Module α [seqLen, vocabularySize] [seqLen, vocabularySize]

Package Gru.Generator as an Spec.Module.

Wiring analogue: GRU language model (nn.GRU + vocabulary projection) producing a sequence of logits.