GRU (spec layer)

TorchLean provides a small GRU specification that is:

• explicit about shapes (so dimension mistakes are caught early),
• explicit about the math (so we can reason about it and differentiate it),
• explicit about which candidate equation is used.

References (math + PyTorch behavior)

• Cho et al., "Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation" (EMNLP 2014): https://aclanthology.org/D14-1179/ (PDF: https://aclanthology.org/D14-1179.pdf)
• Chung et al., "Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling" (2014): https://arxiv.org/abs/1412.3555
• PyTorch GRUCell equations: https://docs.pytorch.org/docs/stable/generated/torch.nn.GRUCell.html
• PyTorch GRU equations: https://docs.pytorch.org/docs/stable/generated/torch.nn.modules.rnn.GRU.html

Notes on parameterization

The GRU equations are often written with separate matrices $W_\bullet$ for the input and $U_\bullet$ for the hidden state. The legacy spec uses a single matrix per gate applied to a concatenated vector $[x_t;h_{t-1}]$ (or $[x_t;r_t\odot h_{t-1}]$ for the candidate). This is the same idea, just packaged in a way that reuses the tensor building blocks already present in the spec layer.

The legacy GRUSpec applies the reset before the hidden-state linear map, as in Cho et al. GRUResetAfterSpec applies it to the recurrent affine output and retains both bias vectors. Use that second specification for PyTorch parameters; the two candidate equations are different functions for general recurrent matrices, so changing tensor layout cannot convert between them.

inductive Spec.GRUConvention : Type

Where the reset gate acts in a GRU candidate.

The original cell resets the hidden vector before multiplying by its recurrent matrix. PyTorch resets the recurrent affine output instead. The distinction matters for a non-diagonal matrix and for a nonzero recurrent candidate bias, so it belongs to the model configuration.

• resetBefore : GRUConvention
• resetAfter : GRUConvention

@[instance_reducible]

instance Spec.instReprGRUConvention : Repr GRUConvention

def Spec.instReprGRUConvention.repr : GRUConvention → ℕ → Std.Format

@[instance_reducible]

instance Spec.instDecidableEqGRUConvention : DecidableEq GRUConvention

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

Reset-after GRU parameters, in PyTorch's packed reset/update/candidate row order.

Rows 0 .. hiddenSize belong to reset, the next block to update, and the last block to the candidate. Input and recurrent weights use [output, input] layout. Both bias vectors remain independent parameters: adding them would lose the candidate's reset-gated recurrent bias and would change how an optimizer updates even the reset and update gates.

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

  PyTorch weight_ih, with reset, update, and candidate rows.

• hiddenWeight : TorchLean.Tensor α [3 * hiddenSize, hiddenSize]

  PyTorch weight_hh, in the same gate order.

• inputBias : TorchLean.Tensor α [3 * hiddenSize]

  PyTorch bias_ih; the candidate part is added outside the reset gate.

• hiddenBias : TorchLean.Tensor α [3 * hiddenSize]

  PyTorch bias_hh; the candidate part is multiplied by the reset gate.

def Spec.GRUResetAfterSpec.ofPyTorch {α : Type} [TorchLean.Storage α] {inputSize hiddenSize : ℕ} (weightIH : TorchLean.Tensor α [3 * hiddenSize, inputSize]) (weightHH : TorchLean.Tensor α [3 * hiddenSize, hiddenSize]) (biasIH biasHH : TorchLean.Tensor α [3 * hiddenSize]) : GRUResetAfterSpec α inputSize hiddenSize

Import a single PyTorch GRU cell's tensors without transposing, merging biases, or changing gates.

The shape indices check the packed row count. A checkpoint's layer and direction selection is the caller's responsibility; these four tensors describe one cell.

def Spec.gruResetAfterCellSpec {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize : ℕ} (gru : GRUResetAfterSpec α inputSize hiddenSize) (input : TorchLean.Tensor α [inputSize]) (prevHidden : TorchLean.Tensor α [hiddenSize]) : TorchLean.Tensor α [hiddenSize]

One reset-after step, with an explicit previous hidden state.

We first compute both packed affine maps, then split their gate blocks. In the candidate, reset * hiddenCandidate includes the recurrent bias because it is already part of that affine map. Keeping this order is what makes copied PyTorch parameters describe the same recurrence.

def Spec.gruResetAfterSequenceSpec {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (gru : GRUResetAfterSpec α inputSize hiddenSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (initialHidden : TorchLean.Tensor α [hiddenSize]) : TorchLean.Tensor α [seqLen, hiddenSize]

Unroll the reset-after cell from the supplied initial state, returning every hidden state.

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

Parameters for a single GRU cell.

This is the original concatenated GRU parameterization, using $[x_t;h_{t-1}]$ (shape inputSize + hiddenSize) for the reset/update gates and $[x_t;r_t\odot h_{t-1}]$ for the candidate gate.

Shapes:

• each gate weight is [hiddenSize, inputSize + hiddenSize],
• each gate bias is [hiddenSize].

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

  Reset-gate weights for $r_t=\operatorname{sigmoid}(W_r[x_t;h_{t-1}]+b_r)$.

• resetBias : TorchLean.Tensor α [hiddenSize]

  Reset-gate bias.

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

  Update-gate weights for $z_t=\operatorname{sigmoid}(W_z[x_t;h_{t-1}]+b_z)$.

• updateBias : TorchLean.Tensor α [hiddenSize]

  Update-gate bias.

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

  Candidate-state weights for $n_t=\tanh(W_n[x_t;r_t\odot h_{t-1}]+b_n)$.

• candidateBias : TorchLean.Tensor α [hiddenSize]

  Candidate-state bias.

def Spec.gruCellSpec {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize : ℕ} (gru : GRUSpec α inputSize hiddenSize) (input : TorchLean.Tensor α [inputSize]) (prevHidden : TorchLean.Tensor α [hiddenSize]) : TorchLean.Tensor α [hiddenSize]

Forward pass for a single GRU cell.

Given input $x_t$ and previous hidden state $h_{t-1}$, compute the next hidden state $h_t$ using the standard GRU equations.

This is not PyTorch's reset-after candidate parameterization; see the module note above.

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

Unroll a GRU over seqLen timesteps (time-major).

This returns the sequence of hidden states $[h_0,\ldots,h_{\mathtt{seqLen}-1}]$. It is a pure spec-level definition of semantics; an efficient runtime is free to implement the same behavior with loops and caching.

The input is time-major and the result contains every hidden state. The candidate semantics remain the Cho-style equations of gruCellSpec.

def Spec.gruCellSpecWithIntermediates {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize : ℕ} (gru : GRUSpec α inputSize hiddenSize) (input : TorchLean.Tensor α [inputSize]) (prevHidden : TorchLean.Tensor α [hiddenSize]) : TorchLean.Tensor α [hiddenSize] × TorchLean.Tensor α [hiddenSize] × TorchLean.Tensor α [hiddenSize] × TorchLean.Tensor α [hiddenSize] × TorchLean.Tensor α [hiddenSize]

GRU cell forward pass that also returns cached intermediates for BPTT.

This computes the same next hidden state as gruCellSpec, but additionally returns:

• resetGate ($r_t$),
• updateGate ($z_t$),
• newCandidate ($n_t$), and
• reset_hidden ($r_t\odot h_{t-1}$).

These are exactly the quantities commonly saved by a reverse-mode implementation (PyTorch-style autograd) to compute gradients efficiently in the backward pass.

def Spec.gruExtractIntermediateValues {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (gru : GRUSpec α inputSize hiddenSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (initialHidden : TorchLean.Tensor α [hiddenSize]) : TorchLean.Tensor α [seqLen, hiddenSize] × TorchLean.Tensor α [seqLen, hiddenSize] × TorchLean.Tensor α [seqLen, hiddenSize] × TorchLean.Tensor α [seqLen, hiddenSize] × TorchLean.Tensor α [seqLen, hiddenSize]

Run a GRU forward pass while collecting the per-timestep intermediates needed for BPTT.

This is the "spec-level" analogue of what frameworks do internally:

• the forward pass produces $h_t$,
• and it also saves gate activations $r_t$, $z_t$, and candidate $n_t$ for the backward pass.

The returned tensors are all time-major (seqLen first) to match the rest of the spec layer.

def Spec.gruBatchedSpec {α : Type} [TorchLean.Storage α] [Context α] {batchSize seqLen inputSize hiddenSize : ℕ} (gru : GRUSpec α inputSize hiddenSize) (inputs : TorchLean.Tensor α [batchSize, seqLen, inputSize]) (initialHidden : TorchLean.Tensor α [batchSize, hiddenSize]) : TorchLean.Tensor α [batchSize, seqLen, hiddenSize]

Batched GRU forward pass (map gruSequenceSpec over the batch dimension).

This is a simple spec-level definition for semantics, not an optimized kernel. It maps the same Cho-style cell over a batch; it is not a torch.nn.GRU checkpoint format.

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

Reference gradient for reset-gate weights via the generic RNN weight-gradient helper.

This uses rnnWeightsDerivSpec on the concatenated inputs/hidden states. It is a convenient building block, but the more explicit BPTT helpers below show the time-unrolled accumulation form.

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

Reference gradient for update-gate weights (via rnnWeightsDerivSpec).

def Spec.gruNewWeightsDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (inputs : TorchLean.Tensor α [seqLen, inputSize]) (resetHiddens gradNew : TorchLean.Tensor α [seqLen, hiddenSize]) : TorchLean.Tensor α [hiddenSize, inputSize + hiddenSize]

Reference gradient for candidate ("new") gate weights (via rnnWeightsDerivSpec).

The second sequence argument satisfies $\mathtt{reset\_hiddens}_t=r_t\odot h_{t-1}$.

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

Bias gradient by summing per-timestep gradients over the time axis.

This is the spec-level analogue of the common "sum across batch/time" reduction used for bias gradients. The seqLen ≠ 0 hypothesis is exactly what makes axis 0 a valid reduction axis.

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

Reset-gate weight gradient by explicit time-unrolled accumulation (BPTT-style).

This computes $$ \sum_t \frac{\partial L}{\partial r_t}\otimes[x_t;h_{t-1}], $$ where $\otimes$ is an outer product.

@[irreducible]

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

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

Update-gate weight gradient by explicit time-unrolled accumulation (BPTT-style).

This computes $$ \sum_t \frac{\partial L}{\partial z_t}\otimes[x_t;h_{t-1}]. $$

@[irreducible]

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

def Spec.gruNewWeightsDerivBpttSpec {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (inputs : TorchLean.Tensor α [seqLen, inputSize]) (resetHiddens gradNewCandidates : TorchLean.Tensor α [seqLen, hiddenSize]) : TorchLean.Tensor α [hiddenSize, inputSize + hiddenSize]

Candidate-gate weight gradient by explicit time-unrolled accumulation (BPTT-style).

This computes $$ \sum_t \frac{\partial L}{\partial n_t}\otimes[x_t;r_t\odot h_{t-1}]. $$

@[irreducible]

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

def Spec.gruCellBackwardFullSpec {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize : ℕ} (gru : GRUSpec α inputSize hiddenSize) (input : TorchLean.Tensor α [inputSize]) (prevHidden gradOutput resetGate updateGate newCandidate : TorchLean.Tensor α [hiddenSize]) : TorchLean.Tensor α [inputSize] × TorchLean.Tensor α [hiddenSize] × TorchLean.Tensor α [hiddenSize, inputSize + hiddenSize] × TorchLean.Tensor α [hiddenSize] × TorchLean.Tensor α [hiddenSize, inputSize + hiddenSize] × TorchLean.Tensor α [hiddenSize] × TorchLean.Tensor α [hiddenSize, inputSize + hiddenSize] × TorchLean.Tensor α [hiddenSize]

Backward (VJP) for a single GRU cell.

Inputs:

• the cell parameters gru,
• the current input $x_t$,
• the previous hidden state $h_{t-1}$,
• an upstream gradient $\partial L/\partial h_t$,
• and the forward intermediates $r_t$, $z_t$, and $n_t$ that a typical BPTT implementation would cache.

Outputs:

• gradients w.r.t. the input and previous hidden state,
• plus gradients for each parameter tensor (weights and biases).

This is written to match the forward equations in gruCellSpec. It is not an optimized kernel; it is a precise spec for what gradients should be.

def Spec.gruSequenceBackwardFullSpec {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (gru : GRUSpec α inputSize hiddenSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (hiddens gradOutputs resetGates updateGates newCandidates : TorchLean.Tensor α [seqLen, hiddenSize]) (initialHidden : TorchLean.Tensor α [hiddenSize] := TorchLean.Tensor.full [hiddenSize] 0) : TorchLean.Tensor α [hiddenSize, inputSize + hiddenSize] × TorchLean.Tensor α [hiddenSize] × TorchLean.Tensor α [hiddenSize, inputSize + hiddenSize] × TorchLean.Tensor α [hiddenSize] × TorchLean.Tensor α [hiddenSize, inputSize + hiddenSize] × TorchLean.Tensor α [hiddenSize] × TorchLean.Tensor α [seqLen, inputSize] × TorchLean.Tensor α [hiddenSize]

Reverse-mode backprop through an unrolled GRU over seqLen steps (BPTT).

This function consumes the same intermediates produced by gruExtractIntermediateValues: per-timestep gate activations and candidates. The backward pass walks time in reverse and accumulates gradients for the Cho-style forward equation.

def Spec.gruSequenceBackwardSpec {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize : ℕ} (gru : GRUSpec α inputSize hiddenSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (hiddens gradOutputs resetGates updateGates newCandidates : TorchLean.Tensor α [seqLen, hiddenSize]) (initialHidden : TorchLean.Tensor α [hiddenSize] := TorchLean.Tensor.full [hiddenSize] 0) : TorchLean.Tensor α [seqLen, inputSize] × TorchLean.Tensor α [hiddenSize]

Return the input-sequence and initial-hidden gradients from gruSequenceBackwardFullSpec.

The full backward pass also returns parameter gradients. This projection records the common contract used by callers that only propagate gradients to the preceding recurrent computation.