RNN/LSTM/GRU module wrappers

The layer specs (NN/Spec/Layers/Rnn.lean, lstm.lean, gru.lean) expose step-level and sequence-level recurrence definitions.

This file wraps the "sequence forward" functions as NNModuleSpecs so recurrent blocks can be composed with other modules in a SpecChain.

Design choices:

• These wrappers are stateless modules: they pick a canonical initial hidden/state (all zeros). This keeps NNModuleSpec simple (pure forward), and makes the semantics explicit. More stateful variants can always be built at the layer-spec level if needed.
• The exported forward returns the full output sequence (not just the final hidden state), matching common encoder usage.

If you think in PyTorch: these are the nn.RNN/nn.LSTM/nn.GRU "return the full output sequence" wrappers, with the initial hidden/state fixed to zeros.

def Spec.RNNModuleSpec {α : Type} [Context α] {seqLen inputSize hiddenSize : ℕ} (rnn : RNNSpec α inputSize hiddenSize) : ModSpec.NNModuleSpec α (Shape.dim seqLen (Shape.dim inputSize Shape.scalar)) (Shape.dim seqLen (Shape.dim hiddenSize Shape.scalar))

RNN sequence wrapper with a zero initial hidden state.

def Spec.LSTMModuleSpec {α : Type} [Context α] {seqLen inputSize hiddenSize : ℕ} (lstm : LSTMSpec α inputSize hiddenSize) : ModSpec.NNModuleSpec α (Shape.dim seqLen (Shape.dim inputSize Shape.scalar)) (Shape.dim seqLen (Shape.dim hiddenSize Shape.scalar))

LSTM sequence wrapper with a zero initial state; returns the output sequence.

def Spec.GRUModuleSpec {α : Type} [Context α] {seqLen inputSize hiddenSize : ℕ} (gru : GRUSpec α inputSize hiddenSize) : ModSpec.NNModuleSpec α (Shape.dim seqLen (Shape.dim inputSize Shape.scalar)) (Shape.dim seqLen (Shape.dim hiddenSize Shape.scalar))

GRU sequence wrapper with a zero initial hidden state; returns the output sequence.

def Spec.BiLSTMModuleSpec {α : Type} [Context α] {seqLen inputSize hiddenSize : ℕ} (forward_lstm backward_lstm : LSTMSpec α inputSize hiddenSize) : ModSpec.NNModuleSpec α (Shape.dim seqLen (Shape.dim inputSize Shape.scalar)) (Shape.dim seqLen (Shape.dim (hiddenSize + hiddenSize) Shape.scalar))

Bidirectional LSTM wrapper (concatenates forward/backward features).

def Spec.RNNCellModuleSpec {α : Type} [Context α] {inputSize hiddenSize : ℕ} (rnn : RNNSpec α inputSize hiddenSize) : ModSpec.NNModuleSpec α (Shape.dim (inputSize + hiddenSize) Shape.scalar) (Shape.dim hiddenSize Shape.scalar)

Wrap rnn_cell_spec as an NNModuleSpec for a single timestep.

Input convention: we take a single vector [x; h] (concatenated input and previous hidden state), so the module is shape-safe and easy to compose.

def Spec.LSTMCellModuleSpec {α : Type} [Context α] {inputSize hiddenSize : ℕ} (lstm : LSTMSpec α inputSize hiddenSize) : ModSpec.NNModuleSpec α (Shape.dim (inputSize + hiddenSize + hiddenSize) Shape.scalar) (Shape.dim (hiddenSize + hiddenSize) Shape.scalar)

Wrap lstm_cell_spec as an NNModuleSpec for a single timestep.

Input convention: a single concatenated vector [x; h; c] (input, previous hidden, previous cell). Output convention: the concatenated new state [h'; c'].

def Spec.GRUCellModuleSpec {α : Type} [Context α] {inputSize hiddenSize : ℕ} (gru : GRUSpec α inputSize hiddenSize) : ModSpec.NNModuleSpec α (Shape.dim (inputSize + hiddenSize) Shape.scalar) (Shape.dim hiddenSize Shape.scalar)

Wrap gru_cell_spec as an NNModuleSpec for a single timestep, using input [x; h].

def Spec.BiRNNModuleSpec {α : Type} [Context α] {seqLen inputSize hiddenSize : ℕ} (forward_rnn backward_rnn : RNNSpec α inputSize hiddenSize) : ModSpec.NNModuleSpec α (Shape.dim seqLen (Shape.dim inputSize Shape.scalar)) (Shape.dim seqLen (Shape.dim (hiddenSize + hiddenSize) Shape.scalar))

Bidirectional RNN wrapper (concatenates forward/backward features).

We run the RNNSpec forward over x, run it again over the reversed sequence, then reverse outputs back and concatenate along the feature axis.