kopitiam_tensor/tensor/lstm.rs
1//! One LSTM timestep: [`Tensor::lstm_cell`].
2//!
3//! This is the recurrence a Tesseract-style LSTM line recognizer runs at
4//! every timestep of a text-line's width. The four gate *pre-activations*
5//! (the summed input + recurrent linear contributions) are computed by the
6//! recognizer with matmuls — [`Tensor::matmul`] or the int8 path in
7//! `tessdata.rs` — and this helper turns them into the next cell state and
8//! hidden output. Keeping the recurrence here (rather than open-coding it in
9//! the recognizer) means the gate math lives next to the `sigmoid`/`tanh`
10//! activations it is built from, and is unit-tested against hand arithmetic
11//! independently of any model.
12//!
13//! The equations are the standard LSTM (Hochreiter & Schmidhuber, 1997;
14//! Gers, Schmidhuber & Cummins, 2000) — clean-room standard math, not a port:
15//!
16//! ```text
17//! i = sigmoid(i_pre) (input gate)
18//! f = sigmoid(f_pre) (forget gate)
19//! o = sigmoid(o_pre) (output gate)
20//! g = tanh(g_pre) (cell candidate)
21//! c' = f * c_prev + i * g (new cell state)
22//! h = o * tanh(c') (new hidden output)
23//! ```
24//!
25//! This is exactly the gate assignment Tesseract's `lstm.cpp` uses (its `CI`
26//! candidate takes `GFunc = Tanh`, its `GI`/`GF1`/`GO` gates take
27//! `FFunc = Logistic`, then `curr_state = f*state + i*g` and
28//! `curr_output = Tanh(state) * o`), but the recurrence itself is textbook —
29//! nothing format- or Tesseract-specific is translated here. Tesseract
30//! additionally clips the cell state to +/-100 before the output squash; that
31//! clamp is a numerical-stability guard, deliberately omitted from this pure
32//! recurrence so the operation stays the plain mathematical LSTM. If a real
33//! model needs it, apply it to the returned [`LstmState::cell`].
34
35use kopitiam_core::{Error, Result};
36
37use super::Tensor;
38
39/// The two things one LSTM timestep produces: the new cell state to carry to
40/// the next timestep, and the hidden output emitted at this one.
41///
42/// A named pair rather than a bare `(Tensor, Tensor)` so a recognizer driving
43/// a sequence cannot silently swap the carried state and the emitted output —
44/// the two have the same shape, so the compiler could not catch a mix-up.
45#[derive(Debug, Clone)]
46pub struct LstmState {
47 /// `c'`, the cell state carried into the next timestep as `cell_prev`.
48 pub cell: Tensor,
49 /// `h`, the hidden output this timestep emits (to a softmax/head, and as
50 /// the recurrent input feeding the next timestep's gate pre-activations).
51 pub hidden: Tensor,
52}
53
54impl Tensor {
55 /// Computes one LSTM timestep from the four gate pre-activations and the
56 /// previous cell state, returning the new [`LstmState`].
57 ///
58 /// Each argument is the *pre-activation* of one gate — the summed input
59 /// and recurrent linear contributions, before any nonlinearity — so this
60 /// helper composes cleanly after the recognizer's gate matmuls. The
61 /// activations (`sigmoid` on the three gates, `tanh` on the candidate and
62 /// on the squashed state) are applied here.
63 ///
64 /// All five tensors must be `f32` and share one shape (the hidden size,
65 /// optionally with leading batch/time dims); the result carries that same
66 /// shape. Broadcasting is intentionally *not* offered: the gates and the
67 /// state of a single LSTM layer are always the same width, so a shape
68 /// disagreement is a bug to surface, not a broadcast to paper over.
69 ///
70 /// # Errors
71 ///
72 /// [`Error::DTypeMismatch`] if any operand is not `f32` (propagated from
73 /// the underlying [`Tensor::sigmoid`]/[`Tensor::tanh`]/[`Tensor::mul`]).
74 /// [`Error::ShapeMismatch`] if the operands' shapes are not all identical.
75 pub fn lstm_cell(
76 input_gate: &Tensor,
77 forget_gate: &Tensor,
78 cell_candidate: &Tensor,
79 output_gate: &Tensor,
80 cell_prev: &Tensor,
81 ) -> Result<LstmState> {
82 let shape = input_gate.shape();
83 for other in [forget_gate, cell_candidate, output_gate, cell_prev] {
84 if other.shape() != shape {
85 return Err(Error::ShapeMismatch {
86 expected: shape.clone(),
87 actual: other.shape().clone(),
88 });
89 }
90 }
91
92 let i = input_gate.sigmoid()?;
93 let f = forget_gate.sigmoid()?;
94 let o = output_gate.sigmoid()?;
95 let g = cell_candidate.tanh()?;
96
97 // c' = f * c_prev + i * g
98 let cell = f.mul(cell_prev)?.add(&i.mul(&g)?)?;
99 // h = o * tanh(c')
100 let hidden = o.mul(&cell.tanh()?)?;
101
102 Ok(LstmState { cell, hidden })
103 }
104}
105
106#[cfg(test)]
107mod tests {
108 use super::*;
109
110 fn assert_close(a: f32, b: f32) {
111 assert!((a - b).abs() < 1e-6, "expected {b}, got {a}");
112 }
113
114 #[test]
115 fn lstm_cell_matches_hand_computation_for_a_single_unit() {
116 // One unit, chosen pre-activations, previous cell state c = 0.5.
117 // i = sigmoid(1.0), f = sigmoid(0.5), o = sigmoid(-0.5), g = tanh(2.0)
118 let i_pre = Tensor::from_f32(vec![1.0], [1]).unwrap();
119 let f_pre = Tensor::from_f32(vec![0.5], [1]).unwrap();
120 let g_pre = Tensor::from_f32(vec![2.0], [1]).unwrap();
121 let o_pre = Tensor::from_f32(vec![-0.5], [1]).unwrap();
122 let c_prev = Tensor::from_f32(vec![0.5], [1]).unwrap();
123
124 let out = Tensor::lstm_cell(&i_pre, &f_pre, &g_pre, &o_pre, &c_prev).unwrap();
125
126 let i = 1.0f32 / (1.0 + (-1.0f32).exp());
127 let f = 1.0f32 / (1.0 + (-0.5f32).exp());
128 let o = 1.0f32 / (1.0 + 0.5f32.exp());
129 let g = 2.0f32.tanh();
130 let c = f * 0.5 + i * g;
131 let h = o * c.tanh();
132
133 assert_close(out.cell.to_vec_f32().unwrap()[0], c);
134 assert_close(out.hidden.to_vec_f32().unwrap()[0], h);
135 }
136
137 #[test]
138 fn lstm_cell_with_zero_input_and_zero_state_stays_at_zero_cell() {
139 // All pre-activations zero, c_prev = 0: g = tanh(0) = 0 and c_prev = 0,
140 // so c' = f*0 + i*0 = 0 exactly. h = o * tanh(0) = 0 too.
141 let z = Tensor::zeros([4]);
142 let out = Tensor::lstm_cell(&z, &z, &z, &z, &z).unwrap();
143 assert_eq!(out.cell.to_vec_f32().unwrap(), vec![0.0; 4]);
144 assert_eq!(out.hidden.to_vec_f32().unwrap(), vec![0.0; 4]);
145 }
146
147 #[test]
148 fn forget_open_and_input_shut_preserves_the_cell_state_exactly() {
149 // f -> 1 (large positive forget pre-activation) and i -> 0 (large
150 // negative input pre-activation) is the LSTM's "carry the memory
151 // unchanged" regime: c' should equal c_prev to within the gates'
152 // saturation error.
153 let big = 30.0f32; // sigmoid(30) ~= 1, sigmoid(-30) ~= 0.
154 let f_pre = Tensor::from_f32(vec![big, big, big], [3]).unwrap();
155 let i_pre = Tensor::from_f32(vec![-big, -big, -big], [3]).unwrap();
156 let g_pre = Tensor::from_f32(vec![5.0, -5.0, 1.0], [3]).unwrap(); // irrelevant: i ~= 0.
157 let o_pre = Tensor::from_f32(vec![0.0, 0.0, 0.0], [3]).unwrap();
158 let c_prev = Tensor::from_f32(vec![0.3, -0.7, 1.2], [3]).unwrap();
159
160 let out = Tensor::lstm_cell(&i_pre, &f_pre, &g_pre, &o_pre, &c_prev).unwrap();
161 let cell = out.cell.to_vec_f32().unwrap();
162 for (got, want) in cell.iter().zip([0.3, -0.7, 1.2]) {
163 assert!((got - want).abs() < 1e-6, "cell state not preserved: got {got}, want {want}");
164 }
165 }
166
167 #[test]
168 fn lstm_cell_preserves_a_multi_dimensional_shape() {
169 let a = Tensor::from_f32(vec![0.1, 0.2, 0.3, 0.4], [2, 2]).unwrap();
170 let out = Tensor::lstm_cell(&a, &a, &a, &a, &a).unwrap();
171 assert_eq!(out.cell.shape().dims(), &[2, 2]);
172 assert_eq!(out.hidden.shape().dims(), &[2, 2]);
173 }
174
175 #[test]
176 fn lstm_cell_rejects_a_shape_disagreement() {
177 let a = Tensor::from_f32(vec![0.0; 3], [3]).unwrap();
178 let b = Tensor::from_f32(vec![0.0; 4], [4]).unwrap();
179 assert!(matches!(
180 Tensor::lstm_cell(&a, &a, &a, &a, &b),
181 Err(Error::ShapeMismatch { .. })
182 ));
183 }
184
185 #[test]
186 fn lstm_cell_rejects_non_f32_input() {
187 let f32t = Tensor::from_f32(vec![0.0; 3], [3]).unwrap();
188 let i32t = Tensor::from_i32(vec![0; 3], [3]).unwrap();
189 assert!(matches!(
190 Tensor::lstm_cell(&i32t, &f32t, &f32t, &f32t, &f32t),
191 Err(Error::DTypeMismatch { .. })
192 ));
193 }
194}