GRU cell, forward (simplified gating)

The problem, from first principles

[04-lstm-cell-forward]'s LSTM solves [03-bptt-vanishing-exploding]'s vanishing-gradient problem with FOUR separate gates and a dedicated cell state c, carried alongside the hidden state h. The GRU (Gated Recurrent Unit, Cho et al., 2014) asks a genuinely useful engineering question: can most of the LSTM's benefit be achieved with LESS machinery? A GRU drops the separate cell state entirely (just one state, h, carried between steps, exactly like a vanilla RNN) and reduces four gates down to TWO: a "reset gate" (how much of the PREVIOUS hidden state to consider when computing a new candidate) and an "update gate" (how much of the OLD hidden state to keep vs. how much of the NEW candidate to blend in). In practice, GRUs are frequently competitive with LSTMs on many tasks, with meaningfully fewer parameters and a simpler implementation, which is part of why they remain a common practical choice even after LSTMs became well established.

The genuinely tricky implementation detail, and the reason this question is more than "the LSTM but with fewer gates": the reset gate r doesn't apply UNIFORMLY to the candidate computation, it multiplies ONLY the hidden-state contribution to the candidate (Whn h + bhn), never the input contribution (Win x + bin). This means the input and hidden pre-activations for the candidate gate have to be computed and kept SEPARATE, rather than summed together first (the way [01-rnn-cell-forward]'s and [04-lstm-cell-forward]'s pre-activations were), since r needs to gate only ONE of the two pieces before they're finally combined.

From theory to code

Implement gru_cell_forward(x, h_prev, weight_ih, weight_hh, bias_ih, bias_hh), matching torch.nn.GRUCell's packing convention: [reset, update, new/candidate], three chunks (not four). Compute gates_ih = x @ weight_ih.T + bias_ih and gates_hh = h_prev @ weight_hh.T + bias_hh SEPARATELY (not summed). Compute r and z by summing their respective gates_ih/gates_hh slices and applying sigmoid. Compute n by applying tanh to gates_ih's candidate slice PLUS r times gates_hh's candidate slice (the reset gate applied ONLY to the hidden contribution). Combine via h_next = (1-z)*n + z*h_prev.

Constraints

• gates_ih and gates_hh must be computed and kept SEPARATE (each its own (batch_size, 3*hidden_size) array), never summed together before the reset gate is applied.
• r and z are each sigmoid of the SUM of their respective gates_ih/gates_hh slices (both contributions summed BEFORE the sigmoid, exactly like [01-rnn-cell-forward]'s combined pre-activation).
• n's hidden-side contribution (gates_hh's candidate slice) must be multiplied by r BEFORE being added to n's input-side contribution (gates_ih's candidate slice); n's input-side contribution is NEVER multiplied by r.
• h_next = (1 - z) * n + z * h_prev: z close to 1 means "keep mostly the OLD hidden state"; z close to 0 means "mostly replace it with the new candidate."

Hints