LSTM cell, forward (gating mechanism)

The problem, from first principles

[03-bptt-vanishing-exploding] demonstrated PRECISELY why vanilla RNNs struggle with long sequences: the same weight_hh matrix gets multiplied into the gradient at every single time step, and that repeated multiplication compounds into vanishing or exploding behavior over enough steps, no matter how carefully weight_hh is initialized. The Long Short-Term Memory cell (Hochreiter & Schmidhuber, 1997) was specifically designed to sidestep this problem, by introducing a SEPARATE "cell state" c, alongside the usual hidden state h, that gets updated PRIMARILY through addition rather than repeated matrix multiplication: c_next = f * c_prev + i * g. That f * c_prev term is an elementwise MULTIPLY by a LEARNED, per-timestep gate value (not a fixed matrix multiplied in every single time regardless of content), and crucially, if the "forget gate" f happens to be close to 1 at a given step (meaning "keep most of the old memory"), the cell state can flow backward through MANY time steps nearly UNCHANGED, gradients don't have to compound through repeated multiplication by the same matrix the way [03-bptt-vanishing-exploding] demonstrated.

An LSTM achieves this by computing FOUR separate gate values at every step, each with its own learned weights, but all packed into the SAME linear transformation for efficiency: an "input gate" (how much of the new candidate information to let in), a "forget gate" (how much of the old cell state to keep vs. discard), a "candidate" (what NEW information might get added), and an "output gate" (how much of the updated cell state to actually expose as the hidden state).

From theory to code

Implement lstm_cell_forward(x, h_prev, c_prev, weight_ih, weight_hh, bias_ih, bias_hh), matching torch.nn.LSTMCell's exact packing convention: weight_ih/weight_hh/bias_ih/bias_hh each contain all FOUR gates' parameters stacked together, in order [input, forget, cell/candidate, output]. Compute the combined pre-activation gates = x @ weight_ih.T + bias_ih + h_prev @ weight_hh.T + bias_hh (shape (batch_size, 4*hidden_size)), split it into four equal hidden_size-wide chunks, apply sigmoid to the input/forget/output gates and tanh to the candidate, then combine via c_next = f*c_prev + i*g and h_next = o*tanh(c_next).

Constraints

• Gates must be extracted from gates in EXACTLY this order: input (i), forget (f), cell/candidate (g), output (o), matching torch.nn.LSTMCell's packing convention.
• i, f, and o use sigmoid (values in (0, 1), interpreted as "how much"); g uses tanh (values in (-1, 1), an actual candidate VALUE, not a gate).
• c_next = f * c_prev + i * g (elementwise multiply, then add), NOT any matrix multiplication.
• h_next = o * tanh(c_next), applying tanh to the UPDATED cell state before gating it with o.

Hints