Gated Recurrent Unit (GRU)#

API#

class warp_nn.modules.layers.GRU(
input_size: int,
hidden_size: int,
*,
num_layers: int = 1,
bidirectional: bool = False,
bias: bool = True,
initialize_parameters: bool = True,
requires_grad: bool = True,
)[source]#

Bases: Module

Apply a multi-layer Gated Recurrent Unit (GRU) to an input sequence.

For each element in the input sequence, each layer computes

\[\begin{split}\begin{array}{ll} r_t = \sigma(W_{ir} \, x_t + b_{ir} + W_{hr} \, h_{t-1} + b_{hr}) \\ z_t = \sigma(W_{iz} \, x_t + b_{iz} + W_{hz} \, h_{t-1} + b_{hz}) \\ n_t = \tanh(W_{in} \, x_t + b_{in} + r_t \odot (W_{hn} \, h_{t-1} + b_{hn})) \\ h_t = (1 - z_t) \odot n_t + z_t \odot h_{t-1} \end{array}\end{split}\]

where \(h_t\) is the hidden state at time \(t\), \(x_t\) is the input at time \(t\) (the output of the previous layer at time \(t\) for the second and subsequent layers), and \(h_{t-1}\) is the hidden state of the layer at time \(t-1\) (or the initial hidden state at time \(0\)), \(r_t\), \(z_t\) and \(n_t\) are the reset, update and new gates, \(\sigma\) is the sigmoid function and \(\odot\) is the element-wise product.

If bidirectional is true, each layer also processes the sequence in reverse order (from the last element to the first one) with its own parameters, and the hidden states of both directions are concatenated (forward direction first) to form the output of the layer.

Learnable parameters (where k is the layer index, and \(D = 2\) if bidirectional is true, otherwise \(D = 1\)):

Name

Shape

Description

\(W_{ir}, W_{iz}, W_{in}\)

weight_ih_l{k}

(3 * hidden_size, input_size) for k = 0, otherwise (3 * hidden_size, D * hidden_size)

Input-to-hidden weights of the k-th layer

\(W_{hr}, W_{hz}, W_{hn}\)

weight_hh_l{k}

(3 * hidden_size, hidden_size)

Hidden-to-hidden weights of the k-th layer

\(b_{ir}, b_{iz}, b_{in}\)

bias_ih_l{k}

(3 * hidden_size, 1)

Input-to-hidden bias of the k-th layer. Only if bias is true

\(b_{hr}, b_{hz}, b_{hn}\)

bias_hh_l{k}

(3 * hidden_size, 1)

Hidden-to-hidden bias of the k-th layer. Only if bias is true

The parameters of the reverse direction (only if bidirectional is true) are named the same, with the _reverse suffix (e.g. weight_ih_l{k}_reverse).

The parameters are initialized from the uniform distribution \(u(-k, k)\) where \(k = \frac{1}{\sqrt{\text{hidden\_size}}}\).

Parameters:

  • input_size – The number of input features.

  • hidden_size – The number of hidden features.

  • num_layers – The number of stacked recurrent layers.

  • bidirectional – Whether each layer processes the sequence in both directions.

  • bias – Whether to include a bias term.

  • initialize_parameters – Whether to initialize the parameters with their default/initial values. If false, the parameters are left as uninitialized memory, and reading them (e.g. during a forward pass) before loading their values (e.g. from a state dictionary) is undefined behavior.

  • requires_grad – Whether the parameters and the cached output arrays of the module require gradients.

Raises:

ValueError – If the number of layers is less than 1.

__call__(
input: array,
hidden: array | None = None,
) → tuple[array, array][source]#

Forward pass of the module.

Parameters:

  • input – The input sequence array, with shape (batch_size, sequence_length, input_size).

  • hidden – The initial hidden state array, with shape (D * num_layers, batch_size, hidden_size). If not given, it defaults to zeros.

Returns:

A tuple of the output sequence array (the hidden states of the last layer for each element in the sequence), with shape (batch_size, sequence_length, D * hidden_size), and the final hidden state array (of each layer and direction), with shape (D * num_layers, batch_size, hidden_size).

Raises:

ValueError – If the shape of the input or hidden state array is not valid.