RNN (spec layer)

Defines a vanilla RNN cell and sequence semantics, along with BPTT-style gradients.

This is the recurrent core that TorchLean builds on:

• a single-step cell (rnnCellSpec),
• an explicit unrolling over time (rnnSequenceSpec),
• and a reverse-time VJP (rnnSequenceBackwardSpec).

PyTorch analogy:

• rnnCellSpec corresponds to torch.nn.RNNCell with nonlinearity="tanh".
• rnnSequenceSpec corresponds to torch.nn.RNN unrolled over seqLen.

References

• Elman, "Finding Structure in Time" (1990): https://crl.ucsd.edu/~elman/Papers/fsit.pdf
• PyTorch RNNCell: https://docs.pytorch.org/docs/stable/generated/torch.nn.RNNCell.html
• PyTorch RNN: https://docs.pytorch.org/docs/stable/generated/torch.nn.RNN.html

Recurrent tensor shapes

Recurrent vectors and matrices use ordinary rank-one and rank-two tensors. Sequences are time-major, with seqLen as the outermost axis, because that layout follows the recursive definitions and proofs directly.

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

RNN cell parameters.

We use a single weight matrix applied to a concatenated vector [x_t; h_{t-1}]:

h_t = tanh(W [x_t; h_{t-1}] + b).

Over exact arithmetic this is equivalent to the common split-parameter form:

h_t = tanh(W_ih x_t + W_hh h_{t-1} + b),

The fused implementation uses one dot-product accumulation across input and hidden coordinates. Two separately rounded products and their sum need not give identical IEEE results.

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

  Combined input-to-hidden and hidden-to-hidden weight matrix.

• bias : TorchLean.Tensor α [hiddenSize]

  Hidden-state bias vector.

def Spec.rnnCellSpec {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize : ℕ} (rnn : RNNSpec α inputSize hiddenSize) (input : TorchLean.Tensor α [inputSize]) (hidden : TorchLean.Tensor α [hiddenSize]) : TorchLean.Tensor α [hiddenSize]

Single RNN cell forward pass.

Math: h_t = tanh(W [x_t; h_{t-1}] + b).

PyTorch analogy: RNNCell(input, hidden) with tanh nonlinearity.

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

Parameter gradients for an RNNSpec cell.

The cell holds a single weight matrix applied to [x_t; h_{t-1}] plus a bias, so this pair is the whole parameter gradient. Encoder and decoder cells use the same record.

PyTorch analogue: (cell.weight_ih.grad, cell.weight_hh.grad) fused into one matrix, plus the bias gradient.

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

  Gradient of the fused input/hidden weight matrix.

• biasGradient : TorchLean.Tensor α [hiddenSize]

  Gradient of the bias.

def Spec.instReprRNNParameterGradients.repr {α✝ : Type} {inst✝ : TorchLean.Storage α✝} {inputSize✝ hiddenSize✝ : ℕ} [Repr α✝] : RNNParameterGradients α✝ inputSize✝ hiddenSize✝ → ℕ → Std.Format

@[instance_reducible]

instance Spec.instReprRNNParameterGradients {α✝ : Type} {inst✝ : TorchLean.Storage α✝} {inputSize✝ hiddenSize✝ : ℕ} [Repr α✝] : Repr (RNNParameterGradients α✝ inputSize✝ hiddenSize✝)

def Spec.RNNParameterGradients.zero {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize : ℕ} : RNNParameterGradients α inputSize hiddenSize

All-zero parameter gradients: the starting point for accumulation over a sequence.

def Spec.RNNParameterGradients.add {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize : ℕ} (left right : RNNParameterGradients α inputSize hiddenSize) : RNNParameterGradients α inputSize hiddenSize

Add two parameter gradient bundles, which is what one BPTT step contributes.

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

Everything one RNN cell step sends backwards.

• parameters : RNNParameterGradients α inputSize hiddenSize

  Gradients for the cell parameters.

• input : TorchLean.Tensor α [inputSize]

  Gradient with respect to the step input x_t.

• previousHidden : TorchLean.Tensor α [hiddenSize]

  Gradient with respect to the incoming hidden state h_{t-1}.

structure Spec.RNNSequenceGradients (α : Type) [TorchLean.Storage α] (seqLen inputSize hiddenSize : ℕ) : Type

Result of backpropagation through time for an RNN: parameter gradients summed over the sequence, one input gradient per timestep, and the gradient for the hidden state fed in at t = 0.

• parameters : RNNParameterGradients α inputSize hiddenSize

  Parameter gradients accumulated over every timestep.

• inputs : TorchLean.Tensor α [seqLen, inputSize]

  Gradient with respect to the input sequence.

• initialHidden : TorchLean.Tensor α [hiddenSize]

  Gradient with respect to the initial hidden state.

def Spec.rnnCellBackwardSpec {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize : ℕ} (rnn : RNNSpec α inputSize hiddenSize) (input : TorchLean.Tensor α [inputSize]) (prevHidden gradHidden : TorchLean.Tensor α [hiddenSize]) : RNNCellGradients α inputSize hiddenSize

Backward/VJP for a single RNN cell.

Inputs:

• x_t, h_{t-1},
• an upstream gradient dL/dh_t.

Outputs:

• an RNNCellGradients record with dL/dx_t, dL/dh_{t-1}, and the parameter gradients.

def Spec.rnnSequenceSpec {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (rnn : RNNSpec α inputSize hiddenSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (initialHidden : TorchLean.Tensor α [hiddenSize]) : TorchLean.Tensor α [seqLen, hiddenSize]

Unroll an RNN over seqLen steps (time-major).

Returns the sequence of hidden states [h_0, ..., h_{seqLen-1}].

def Spec.rnnSequenceBackwardSpec {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (rnn : RNNSpec α inputSize hiddenSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (initialHidden : TorchLean.Tensor α [hiddenSize]) (hiddens gradHiddens : TorchLean.Tensor α [seqLen, hiddenSize]) : RNNSequenceGradients α seqLen inputSize hiddenSize

Full BPTT backward pass through an RNN sequence.

This is the spec-level version of what PyTorch autograd computes for nn.RNN when unrolled:

• we walk time in reverse,
• accumulate parameter gradients,
• and compute gradients for each input step plus the initial hidden state.

initialHidden may be nonzero; it supplies the previous state at the first timestep. For an empty sequence, both parameter gradients and the gradient of that initial state are zero.

Diagram: forward unroll + BPTT (vanilla RNN)

One step (forward):

x_t        h_{t-1}
 |            |
 +---- concat ----+
                 |
             z_t = W · [x_t; h_{t-1}] + b
                 |
             h_t = tanh(z_t)

Unrolled over time (forward):

h_-1 = h0

x0 -> [cell] -> h0 -> [cell] -> h1 -> ... -> [cell] -> h_{T-1}
        ^          ^                       ^
      uses h_-1  uses h0                 uses h_{T-2}

Backprop through time (reverse):

At each time step we combine two sources of gradient for h_t:

• the gradient coming from the loss that touches h_t directly (gradHiddens[t]),
• plus the gradient flowing "from the future" through the recurrence (dHidden_next).

Then we push total_grad through the single-step VJP (rnnCellBackwardSpec), producing:

• dInput_t and dHidden_prev,
• and parameter gradients which are accumulated across time.