Module ml::rnn
Recurrent cells: RNN, GRU, LSTM, xLSTM.
All vectors are row tensors: x_t is [1, inputSize], h_t/c_t are [1, hiddenSize]. Weight matrices are laid out so a plain tensor_matmul(x_t, W) gives the right shape directly (i.e. W is [inputSize, hiddenSize] or [hiddenSize, hiddenSize], the transpose of PyTorch's own [hiddenSize, inputSize] convention — chosen so this library's existing row-vector Tensor shape needs no extra transpose per step). Forward-only, like the rest of std/ml (see attention.h).
Structs
| RNNCell | h_t = tanh(x_t @ W_ih + h_{t-1} @ W_hh + b_ih + b_hh) |
| GRUCell | r_t = sigmoid(x_t@W_ir + h@W_hr + b_ir + b_hr) reset gate z_t = sigmoid(x_t@W_iz + h@W_hz + b_iz + b_hz) update gate n_t = tanh(x_t@W_in + b_in + r_t * (h@W_hn + b_hn)) candidate h_t = (1 - z_t) * n_t + z_t * h_{t-1} |
| LSTMCell | i_t = sigmoid(x_t@W_ii + h@W_hi + b_ii + b_hi) input gate f_t = sigmoid(x_t@W_if + h@W_hf + b_if + b_hf) forget gate g_t = tanh(x_t@W_ig + h@W_hg + b_ig + b_hg) cell candidate o_t = sigmoid(x_t@W_io + h@W_ho + b_io + b_ho) output gate c_t = f_t * c_{t-1} + i_t * g_t h_t = o_t * tanh(c_t) |
| XLSTMCell | The scalar-memory sLSTM cell — xLSTM's headline idea is exponential gating (replacing LSTM's sigmoid input/forget gates with unbounded exp(), which can represent much larger gate values) made numerically stable with a running log-domain stabilizer state 'm_t', plus a normalizer state 'n_t' that keeps the exponential gates' magnitude in check without clipping. This is the scalar-memory (sLSTM) member of the xLSTM family specifically — mLSTM (matrix-memory, fully parallelizable) is a separate, larger cell not implemented here. z_t = tanh(x_t@Wz + h@Rz + bz) cell input i_raw = x_t@Wi + h@Ri + bi; f_raw = x_t@Wf + h@Rf + bf (pre-activation logits) o_t = sigmoid(x_t@Wo + h@Ro + bo) m_t = max(f_raw + m_{t-1}, i_raw) stabilizer (log-domain) i_t = exp(i_raw - m_t); f_t = exp(f_raw + m_{t-1} - m_t) c_t = f_t * c_{t-1} + i_t * z_t n_t = f_t * n_{t-1} + i_t h_t = o_t * (c_t / n_t) |
| XLSTMState |