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

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.

Both conventions evaluate the complementary update gate as sigmoid (-preUpdate), rather than subtracting its rounded output from one. The candidate and its gradient can remain nonzero even after the update gate itself has rounded to one.

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

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

@[instance_reducible]

instance Spec.instReprGRUConvention : Repr GRUConvention

@[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.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 returns the next hidden state together with:

• 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.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 the first component of gruCellSpecWithIntermediates, so the two cannot drift apart.

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.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.gruCellBackwardSpec {α : 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.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 α [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.