LSTM (spec layer)

TorchLean provides a small LSTM specification that is:

• explicit about shapes (so common dimension mistakes are caught early),
• explicit about the gate math (so gradients are inspectable and proofs can refer to the equations),
• close in spirit to the way PyTorch documents nn.LSTMCell / nn.LSTM.

References (math + PyTorch behavior)

• Hochreiter, Schmidhuber, "Long Short-Term Memory" (Neural Computation, 1997). Free PDF: http://www.bioinf.jku.at/publications/older/2604.pdf
• PyTorch LSTMCell equations: https://docs.pytorch.org/docs/stable/generated/torch.nn.LSTMCell.html
• PyTorch LSTM equations: https://docs.pytorch.org/docs/stable/generated/torch.nn.LSTM.html

Notes on parameterization

Many libraries expose two matrices per gate (W_ih and W_hh) and add them. In this spec we use a single matrix applied to a concatenated vector [x_t; h_{t-1}]. It's the same computation, just packaged to reuse TorchLean's tensor building blocks.

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

Parameters for an LSTM cell, with one (hiddenSize × (inputSize + hiddenSize)) matrix per gate.

This corresponds to the usual (W_ih, W_hh) parameterization in libraries like PyTorch, but we package it as a single matrix applied to [x_t; h_{t-1}] to reuse TorchLean's tensor building blocks.

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

  Forget-gate weights for f_t = sigmoid(W_f [x_t; h_{t-1}] + b_f).

• forgetBias : TorchLean.Tensor α [hiddenSize]

  Forget-gate bias.

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

  Input-gate weights for i_t = sigmoid(W_i [x_t; h_{t-1}] + b_i).

• inputBias : TorchLean.Tensor α [hiddenSize]

  Input-gate bias.

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

  Candidate/cell-proposal weights for g_t = tanh(W_g [x_t; h_{t-1}] + b_g).

• candidateBias : TorchLean.Tensor α [hiddenSize]

  Candidate/cell-proposal bias.

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

  Output-gate weights for o_t = sigmoid(W_o [x_t; h_{t-1}] + b_o).

• outputBias : TorchLean.Tensor α [hiddenSize]

  Output-gate bias.

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

LSTM recurrent state: hidden vector h_t and cell vector c_t.

• hidden : TorchLean.Tensor α [hiddenSize]

  Exposed hidden state h_t.

• cell : TorchLean.Tensor α [hiddenSize]

  Internal memory/cell state c_t.

def Spec.lstmCellSpec {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize : ℕ} (lstm : LSTMSpec α inputSize hiddenSize) (input : TorchLean.Tensor α [inputSize]) (prevState : LSTMState α hiddenSize) : LSTMState α hiddenSize

One LSTM cell step: update (h_{t-1}, c_{t-1}) given x_t and parameters.

def Spec.lstmSequenceSpec {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (lstm : LSTMSpec α inputSize hiddenSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (initialState : LSTMState α hiddenSize) : TorchLean.Tensor α [seqLen, hiddenSize] × LSTMState α hiddenSize

Run an LSTM cell over a length-seqLen input sequence, returning outputs and final state.

def Spec.lstmBatchedSpec {α : Type} [TorchLean.Storage α] [Context α] {batchSize seqLen inputSize hiddenSize : ℕ} (lstm : LSTMSpec α inputSize hiddenSize) (inputs : TorchLean.Tensor α [batchSize, seqLen, inputSize]) (initialHiddens : TorchLean.Tensor α [batchSize, hiddenSize]) : TorchLean.Tensor α [batchSize, seqLen, hiddenSize] × TorchLean.Tensor α [batchSize, hiddenSize]

Batched wrapper around lstmSequenceSpec (runs one sequence per batch element).

def Spec.lstmCellSpecWithIntermediates {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize : ℕ} (lstm : LSTMSpec α inputSize hiddenSize) (input : TorchLean.Tensor α [inputSize]) (prevState : LSTMState α hiddenSize) : LSTMState α hiddenSize × TorchLean.Tensor α [hiddenSize] × TorchLean.Tensor α [hiddenSize] × TorchLean.Tensor α [hiddenSize] × TorchLean.Tensor α [hiddenSize]

Forward pass for one LSTM cell that also returns the gate activations.

This is the spec analogue of the "saved tensors" that a runtime will keep for backward.

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

Gate-wise parameter gradients for an LSTM cell.

LSTMSpec keeps one weight matrix per gate, each applied to the concatenation [x_t; h_{t-1}], so every weight gradient has shape [hiddenSize, inputSize + hiddenSize] and every bias gradient has shape [hiddenSize]. That uniformity is the reason for this record: the eight tensors used to travel as a positional tuple, where four identically shaped weight/bias pairs meant a swapped gate was invisible to the type checker, and the BPTT loop below had to thread them through a nine-element accumulator. Names cost nothing and catch that class of mistake at the call site.

PyTorch analogue: the .grad fields of nn.LSTMCell.weight_ih, weight_hh and their biases, with the input and hidden blocks kept in one matrix here rather than two.

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

  Gradient of the forget-gate weight matrix.

• forgetBias : TorchLean.Tensor α [hiddenSize]

  Gradient of the forget-gate bias.

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

  Gradient of the input-gate weight matrix.

• inputBias : TorchLean.Tensor α [hiddenSize]

  Gradient of the input-gate bias.

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

  Gradient of the candidate-state weight matrix.

• candidateBias : TorchLean.Tensor α [hiddenSize]

  Gradient of the candidate-state bias.

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

  Gradient of the output-gate weight matrix.

• outputBias : TorchLean.Tensor α [hiddenSize]

  Gradient of the output-gate bias.

@[instance_reducible]

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

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

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

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

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

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

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

Everything one LSTM cell step sends backwards: gate parameter gradients, the input gradient, and the gradient for the state that arrived from the previous step.

• gates : LSTMGateGradients α inputSize hiddenSize

  Gradients for the four gate parameter blocks.

• input : TorchLean.Tensor α [inputSize]

  Gradient with respect to the step input x_t.

• previousState : LSTMState α hiddenSize

  Gradient with respect to the incoming state (h_{t-1}, c_{t-1}).

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

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

• gates : LSTMGateGradients α inputSize hiddenSize

  Gate parameter gradients accumulated over every timestep.

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

  Gradient with respect to the input sequence.

• initialState : LSTMState α hiddenSize

  Gradient with respect to the initial state.

def Spec.lstmCellBackwardSpec {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize : ℕ} (lstm : LSTMSpec α inputSize hiddenSize) (input : TorchLean.Tensor α [inputSize]) (prevState state : LSTMState α hiddenSize) (forgetGate inputGate candidate outputGate gradHidden gradCell : TorchLean.Tensor α [hiddenSize]) : LSTMCellGradients α inputSize hiddenSize

Backward pass (VJP) for a single LSTM cell.

Inputs:

• parameters lstm,
• inputs x_t, previous state (h_{t-1}, c_{t-1}), and current state (h_t, c_t),
• the gate activations from the forward pass,
• upstream gradients for both h_t and c_t.

Outputs, as an LSTMCellGradients record: gradients w.r.t. x_t and the previous state, plus one gradient per gate parameter tensor.

This is the quantity computed by PyTorch autograd for an nn.LSTMCell unrolled in time.

def Spec.lstmSequenceBackwardSpec {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (lstm : LSTMSpec α inputSize hiddenSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (initialState : LSTMState α hiddenSize) (gradHiddens : TorchLean.Tensor α [seqLen, hiddenSize]) : LSTMSequenceGradients α seqLen inputSize hiddenSize

Backprop through time (BPTT) for the whole sequence.

This function recomputes and stores the forward intermediates (gates and states) internally, then walks time backward accumulating parameter gradients and input gradients. This matches the usual PyTorch training structure, with the save-vs-recompute choice made explicit.