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).

This is equivalent to the common split-parameter form:

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

just packaged to reuse the same tensor primitives elsewhere in TorchLean.

• 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. The seq2seq baseline in NN/Spec/Models/Seq2seq.lean used to declare its own identical copy of this record; sharing one means an encoder gradient and a decoder gradient have the same type.

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 hidden gradHidden : TorchLean.Tensor α [hiddenSize]) : RNNCellGradients α inputSize hiddenSize

Backward/VJP for a single RNN cell.

Inputs:

• x_t, h_{t-1},
• the cached forward output h_t (so we can write tanh' in terms of h_t),
• 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.rnnBatchedSpec {α : Type} [TorchLean.Storage α] [Context α] {batchSize seqLen inputSize hiddenSize : ℕ} (rnn : RNNSpec α inputSize hiddenSize) (inputs : TorchLean.Tensor α [batchSize, seqLen, inputSize]) (initialHidden : TorchLean.Tensor α [batchSize, hiddenSize]) : TorchLean.Tensor α [batchSize, seqLen, hiddenSize]

Batched RNN forward pass (maps rnnSequenceSpec over the batch dimension).

def Spec.rnnWeightsDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (inputs : TorchLean.Tensor α [seqLen, inputSize]) (hiddens gradOutputs : TorchLean.Tensor α [seqLen, hiddenSize]) : TorchLean.Tensor α [hiddenSize, inputSize + hiddenSize]

Gradient w.r.t. weights from a full unroll, given per-step preactivation gradients.

This helper is for analyses that already have preactivation gradients. It assumes:

• the initial hidden state is 0, and
• gradOutputs[t] is already dL/dz_t (preactivation gradient).

For end-to-end BPTT from dL/dh_t, prefer rnnSequenceBackwardSpec.

@[irreducible]

def Spec.rnnWeightsDerivSpec.accumulate_grads {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (inputs : TorchLean.Tensor α [seqLen, inputSize]) (hiddens gradOutputs : TorchLean.Tensor α [seqLen, hiddenSize]) (t : ℕ) (acc : TorchLean.Tensor α [hiddenSize, inputSize + hiddenSize]) : TorchLean.Tensor α [hiddenSize, inputSize + hiddenSize]

def Spec.rnnBiasDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {seqLen hiddenSize : ℕ} (gradOutputs : TorchLean.Tensor α [seqLen, hiddenSize]) (h : seqLen ≠ 0) : TorchLean.Tensor α [hiddenSize]

Gradient w.r.t. bias from per-step preactivation gradients.

This is sum_t dL/dz_t over the sequence dimension.

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.

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.