GRU Cell Step

Update a recurrent hidden state

A gated recurrent unit (GRU) combines a new input with a previous hidden state. Implement one step using explicit Linen parameters so you can see how each gate affects the result. This function processes one input vector, not an entire sequence or batch.

Let x have shape (D,) and h_prev have shape (H,). Concatenate them in that order and compute:

xh        = concatenate([x, h_prev])
r         = sigmoid(xh @ W_r + b_r)
z         = sigmoid(xh @ W_z + b_z)
xrh       = concatenate([x, r * h_prev])
candidate = tanh(xrh @ W_h + b_h)
h_new     = (1 - z) * h_prev + z * candidate

Each weight matrix has shape (D + H, H) and each bias has shape (H,). The matrix multiplies from the right, matching JAX's row-vector convention. Gate products and the final blend are elementwise.

Read the gates using this exercise's convention

The update gate z controls the contribution of the candidate. Values near zero mostly retain the previous state; values near one mostly use the candidate. Exact zero and one describe the limiting behavior of the sigmoid, not a promise of perfect long-term memory.

Other GRU definitions use z for the amount of the old state to keep. For example, PyTorch's GRUCell weights the old state by z. That convention is valid, but it does not match this exercise's seeded parameters. Reversing the blend here changes the answer unless you also transform the update gate consistently.

The reset gate r scales the old hidden state before the candidate's linear projection. With a reset value near zero, that coordinate of the old state contributes little to the candidate. The old state still contributes directly through the final blend, and it also influences how the gates are computed.

Reset before projecting

Split the candidate matrix into its first D rows and last H rows:

xrh @ W_h = x @ W_h[:D] + (r * h_prev) @ W_h[D:]

Do not replace the second term with r * (h_prev @ W_h[D:]). A matrix can mix hidden coordinates, so scaling before and after that mixing generally gives different results. This is another reason a built-in GRU is not a drop-in replacement for these equations.

Concatenating input and hidden state is one way to store the weights. Separate input and hidden matrices are equivalent when they are slices of this same matrix; independently initializing them would not reproduce the seeded weights.

A step with known weights

Use zero gate weights and biases so both gates equal 0.5. Choose candidate weights that make the reset operation easy to see:

import jax
import jax.numpy as jnp

x = jnp.array([1.])
h_prev = jnp.array([2., -2.])
gate_weight = jnp.zeros((3, 2))
gate_bias = jnp.zeros(2)
W_h = jnp.array([[1., 0.], [0., 1.], [1., 1.]])
xh = jnp.concatenate([x, h_prev])
r = jax.nn.sigmoid(xh @ gate_weight + gate_bias)
z = jax.nn.sigmoid(xh @ gate_weight + gate_bias)
xrh = jnp.concatenate([x, r * h_prev])
candidate = jnp.tanh(xrh @ W_h)
h_new = (1 - z) * h_prev + z * candidate

assert jnp.array_equal(xrh, jnp.array([1., 1., -1.]))
assert jnp.array_equal(candidate, jnp.zeros(2))
assert jnp.array_equal(h_new, jnp.array([1., -1.]))

Here gate_weight and gate_bias stand in for two independently stored gate parameter pairs. The exercise must create all six named arrays even though these example values happen to be equal.