Recurrent Models

This file builds higher-level recurrent architectures by composing module specifications:

• sequence‑to‑sequence: RNN over inputs + per‑step linear projection,
• many‑to‑one classification: RNN + classifier head on the final hidden state,
• bidirectional variants (where supported by module specs).

The cell dynamics live in NN.Spec.Layers.Rnn; this module defines model-level compositions and reference forward and backward functions.

def Spec.Rnn.sequence {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (rnnSpec : RNNSpec α inputSize hiddenSize) (linearSpec : LinearSpec α hiddenSize outputSize) : Module.Chain α [seqLen, inputSize] [seqLen, outputSize]

A simple sequence-to-sequence RNN "model wiring" expressed as a Spec.Module.Chain.

This composes Spec.Module.rnn with a linear projection applied independently at each timestep, so the overall model maps:

• input shape: [seqLen, inputSize]
• output shape: [seqLen, outputSize]

PyTorch analogue: applying nn.RNN (or a custom recurrent cell) followed by a nn.linear at each time step.

def Spec.Rnn.classifier {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize numClasses : ℕ} (rnnSpec : RNNSpec α inputSize hiddenSize) (classifierHead : LinearSpec α hiddenSize numClasses) (h : seqLen ≠ 0) : Module.Chain α [seqLen, inputSize] [numClasses]

A many-to-one RNN classifier expressed as a Spec.Module.Chain.

This runs an RNN over the input sequence and then applies a linear classifier head to the final hidden state.

PyTorch analogue: nn.RNN (or nn.GRU/nn.LSTM) feeding a nn.linear head, taking the last hidden/output.

def Spec.Rnn.stacked {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (layers : RecurrentStack (RNNSpec α) 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.Rnn.languageModel {α : Type} [TorchLean.Storage α] [Context α] {seqLen vocabularySize hiddenSize : ℕ} (embeddingSpec : LinearSpec α vocabularySize hiddenSize) (rnnSpec : RNNSpec α hiddenSize hiddenSize) (outputSpec : LinearSpec α hiddenSize vocabularySize) : Module.Chain α [seqLen, vocabularySize] [seqLen, vocabularySize]

A simple RNN language model spec: "embedding" linear map, RNN core, and output projection.

This file treats embedding/projection as LinearSpecs. A common spec-level usage is that tokens are one-hot vectors of length vocabularySize, so the embedding is just a matrix multiply.

PyTorch analogue: nn.Embedding (conceptually) + nn.RNN + nn.linear vocabulary projection.

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

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

This is a "record of specs" representation, as opposed to the Spec.Module.Chain representation used above.

• rnn : RNNSpec α inputSize hiddenSize

  Recurrent cell parameters.

• outputLayer : LinearSpec α hiddenSize outputSize

  Linear output projection.

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

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

Both halves reuse the gradient records of the layers they belong to, so a caller that already knows how to read Spec.RNNParameterGradients needs to learn nothing new here.

• rnn : RNNParameterGradients α inputSize hiddenSize

  Gradients of the recurrent cell parameters.

• output : LinearParameterGradients α hiddenSize outputSize

  Gradients of the linear output projection.

structure Spec.Rnn.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 (RNNSpec α) inputSize hiddenSize

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

• outputLayer : LinearSpec α hiddenSize outputSize

  Linear output projection.

def Spec.Rnn.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.Rnn.Classifier (α : Type) [TorchLean.Storage α] (inputSize hiddenSize numClasses : ℕ) : Type

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

The classifier head is a linear layer applied to the final hidden state.

• rnn : RNNSpec α inputSize hiddenSize

  Recurrent cell parameters.

• classifier : LinearSpec α hiddenSize numClasses

  Linear classifier head.

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

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

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

• embedding : LinearSpec α vocabularySize hiddenSize

  Token projection used by this one-hot specification.

• rnn : RNNSpec α hiddenSize hiddenSize

  Recurrent cell parameters.

• outputProjection : LinearSpec α hiddenSize vocabularySize

  Projection from hidden states to vocabulary logits.

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

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

The output head consumes the concatenation of forward and backward hidden states.

• forwardRnn : RNNSpec α inputSize hiddenSize

  Recurrent cell for the original sequence order.

• backwardRnn : RNNSpec α inputSize hiddenSize

  Recurrent cell for the reversed sequence order.

• outputLayer : LinearSpec α (hiddenSize + hiddenSize) outputSize

  Projection from concatenated forward and backward states.

def Spec.Rnn.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 pass for Rnn.Model.

Given an input vector and current hidden state, compute the output and next state using the RNN cell and the linear head.

def Spec.Rnn.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 pass for Rnn.Model.

Runs the recurrent cell over the full sequence, applies the output layer at each time step, and returns both the per-step outputs and the final hidden state. Empty sequences preserve the initial hidden state.

def Spec.Rnn.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 an Rnn.Classifier (many-to-one).

This runs the recurrent core over the input sequence and feeds the last hidden state to the classifier head.

def Spec.Rnn.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 an Rnn.Generator (many-to-many).

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

def Spec.Rnn.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 RNN model.

This runs a forward RNN on the sequence, a backward RNN on the reversed sequence, concatenates the two state streams per time step, and applies the output head.

def Spec.Rnn.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]) (gradOutputs : TorchLean.Tensor α [seqLen, outputSize]) : Grads α inputSize hiddenSize outputSize × TorchLean.Tensor α [seqLen, inputSize]

Backward pass for Rnn.Model over a full sequence.

Returns parameter gradients together with the gradient of the input sequence. Empty sequences produce zero parameter gradients and an empty input gradient.

hiddenStates must come from rnnSequenceSpec with a zero initial hidden state, as used by Model.toModule.

This is a spec-level reference implementation; performance is not a goal here.

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

Package an Rnn.Model as a shape-indexed module.

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

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

Package an Rnn.Classifier as an Spec.Module.

This plugs the classifier into the common module pipeline and records a PyTorch-oriented summary.