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 Spec.Modules so recurrent blocks can be composed with other modules in a Spec.Module.Chain.

Design choices:

• These wrappers are stateless modules: they pick a canonical initial hidden/state (all zeros). Spec.Module remains a pure forward; more stateful variants can be built at the layer-spec level if needed.
• The exported forward returns the full output sequence, including every intermediate 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.Module.rnn {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (rnn : RNNSpec α inputSize hiddenSize) : Module α [seqLen, inputSize] [seqLen, hiddenSize]

RNN sequence wrapper with a zero initial hidden state.

def Spec.Module.lstm {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (lstm : LSTMSpec α inputSize hiddenSize) : Module α [seqLen, inputSize] [seqLen, hiddenSize]

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

def Spec.Module.gru {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (gru : GRUSpec α inputSize hiddenSize) : Module α [seqLen, inputSize] [seqLen, hiddenSize]

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

def Spec.Module.bidirectionalLstm {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (forwardLstm backwardLstm : LSTMSpec α inputSize hiddenSize) : Module α [seqLen, inputSize] [seqLen, hiddenSize + hiddenSize]

Bidirectional LSTM wrapper (concatenates forward/backward features).

def Spec.Module.rnnCell {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize : ℕ} (rnn : RNNSpec α inputSize hiddenSize) : Module α [inputSize + hiddenSize] [hiddenSize]

Wrap rnnCellSpec as an Spec.Module 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.Module.lstmCell {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize : ℕ} (lstm : LSTMSpec α inputSize hiddenSize) : Module α [inputSize + hiddenSize + hiddenSize] [hiddenSize + hiddenSize]

Wrap lstmCellSpec as an Spec.Module 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.Module.gruCell {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize : ℕ} (gru : GRUSpec α inputSize hiddenSize) : Module α [inputSize + hiddenSize] [hiddenSize]

Wrap gruCellSpec as an Spec.Module for a single timestep, using input [x; h].

def Spec.Module.bidirectionalRnn {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (forwardRnn backwardRnn : RNNSpec α inputSize hiddenSize) : Module α [seqLen, inputSize] [seqLen, hiddenSize + hiddenSize]

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.