burn_nn/modules/rnn/
lstm.rs

1use burn_core as burn;
2
3use crate::GateController;
4use burn::config::Config;
5use burn::module::Module;
6use burn::module::{Content, DisplaySettings, Initializer, ModuleDisplay};
7use burn::tensor::Tensor;
8use burn::tensor::activation;
9use burn::tensor::backend::Backend;
10
11/// A LstmState is used to store cell state and hidden state in LSTM.
12pub struct LstmState<B: Backend, const D: usize> {
13    /// The cell state.
14    pub cell: Tensor<B, D>,
15    /// The hidden state.
16    pub hidden: Tensor<B, D>,
17}
18
19impl<B: Backend, const D: usize> LstmState<B, D> {
20    /// Initialize a new [LSTM State](LstmState).
21    pub fn new(cell: Tensor<B, D>, hidden: Tensor<B, D>) -> Self {
22        Self { cell, hidden }
23    }
24}
25
26/// Configuration to create a [Lstm](Lstm) module using the [init function](LstmConfig::init).
27#[derive(Config, Debug)]
28pub struct LstmConfig {
29    /// The size of the input features.
30    pub d_input: usize,
31    /// The size of the hidden state.
32    pub d_hidden: usize,
33    /// If a bias should be applied during the Lstm transformation.
34    pub bias: bool,
35    /// Lstm initializer
36    #[config(default = "Initializer::XavierNormal{gain:1.0}")]
37    pub initializer: Initializer,
38}
39
40/// The Lstm module. This implementation is for a unidirectional, stateless, Lstm.
41///
42/// Introduced in the paper: [Long Short-Term Memory](https://www.researchgate.net/publication/13853244).
43///
44/// Should be created with [LstmConfig].
45#[derive(Module, Debug)]
46#[module(custom_display)]
47pub struct Lstm<B: Backend> {
48    /// The input gate regulates which information to update and store in the cell state at each time step.
49    pub input_gate: GateController<B>,
50    /// The forget gate is used to control which information to discard or keep in the memory cell at each time step.
51    pub forget_gate: GateController<B>,
52    /// The output gate determines which information from the cell state to output at each time step.
53    pub output_gate: GateController<B>,
54    /// The cell gate is used to compute the cell state that stores and carries information through time.
55    pub cell_gate: GateController<B>,
56    /// The hidden state of the LSTM.
57    pub d_hidden: usize,
58}
59
60impl<B: Backend> ModuleDisplay for Lstm<B> {
61    fn custom_settings(&self) -> Option<DisplaySettings> {
62        DisplaySettings::new()
63            .with_new_line_after_attribute(false)
64            .optional()
65    }
66
67    fn custom_content(&self, content: Content) -> Option<Content> {
68        let [d_input, _] = self.input_gate.input_transform.weight.shape().dims();
69        let bias = self.input_gate.input_transform.bias.is_some();
70
71        content
72            .add("d_input", &d_input)
73            .add("d_hidden", &self.d_hidden)
74            .add("bias", &bias)
75            .optional()
76    }
77}
78
79impl LstmConfig {
80    /// Initialize a new [lstm](Lstm) module.
81    pub fn init<B: Backend>(&self, device: &B::Device) -> Lstm<B> {
82        let d_output = self.d_hidden;
83
84        let new_gate = || {
85            GateController::new(
86                self.d_input,
87                d_output,
88                self.bias,
89                self.initializer.clone(),
90                device,
91            )
92        };
93
94        Lstm {
95            input_gate: new_gate(),
96            forget_gate: new_gate(),
97            output_gate: new_gate(),
98            cell_gate: new_gate(),
99            d_hidden: self.d_hidden,
100        }
101    }
102}
103
104impl<B: Backend> Lstm<B> {
105    /// Applies the forward pass on the input tensor. This LSTM implementation
106    /// returns the state for each element in a sequence (i.e., across seq_length) and a final state.
107    ///
108    /// ## Parameters:
109    /// - batched_input: The input tensor of shape `[batch_size, sequence_length, input_size]`.
110    /// - state: An optional `LstmState` representing the initial cell state and hidden state.
111    ///   Each state tensor has shape `[batch_size, hidden_size]`.
112    ///   If no initial state is provided, these tensors are initialized to zeros.
113    ///
114    /// ## Returns:
115    /// - output: A tensor represents the output features of LSTM. Shape: `[batch_size, sequence_length, hidden_size]`
116    /// - state: A `LstmState` represents the final states. Both `state.cell` and `state.hidden` have the shape
117    ///   `[batch_size, hidden_size]`.
118    pub fn forward(
119        &self,
120        batched_input: Tensor<B, 3>,
121        state: Option<LstmState<B, 2>>,
122    ) -> (Tensor<B, 3>, LstmState<B, 2>) {
123        let device = batched_input.device();
124        let [batch_size, seq_length, _] = batched_input.dims();
125
126        self.forward_iter(
127            batched_input.iter_dim(1).zip(0..seq_length),
128            state,
129            batch_size,
130            seq_length,
131            &device,
132        )
133    }
134
135    fn forward_iter<I: Iterator<Item = (Tensor<B, 3>, usize)>>(
136        &self,
137        input_timestep_iter: I,
138        state: Option<LstmState<B, 2>>,
139        batch_size: usize,
140        seq_length: usize,
141        device: &B::Device,
142    ) -> (Tensor<B, 3>, LstmState<B, 2>) {
143        let mut batched_hidden_state =
144            Tensor::empty([batch_size, seq_length, self.d_hidden], device);
145
146        let (mut cell_state, mut hidden_state) = match state {
147            Some(state) => (state.cell, state.hidden),
148            None => (
149                Tensor::zeros([batch_size, self.d_hidden], device),
150                Tensor::zeros([batch_size, self.d_hidden], device),
151            ),
152        };
153
154        for (input_t, t) in input_timestep_iter {
155            let input_t = input_t.squeeze_dim(1);
156            // f(orget)g(ate) tensors
157            let biased_fg_input_sum = self
158                .forget_gate
159                .gate_product(input_t.clone(), hidden_state.clone());
160            let forget_values = activation::sigmoid(biased_fg_input_sum); // to multiply with cell state
161
162            // i(nput)g(ate) tensors
163            let biased_ig_input_sum = self
164                .input_gate
165                .gate_product(input_t.clone(), hidden_state.clone());
166            let add_values = activation::sigmoid(biased_ig_input_sum);
167
168            // o(output)g(ate) tensors
169            let biased_og_input_sum = self
170                .output_gate
171                .gate_product(input_t.clone(), hidden_state.clone());
172            let output_values = activation::sigmoid(biased_og_input_sum);
173
174            // c(ell)g(ate) tensors
175            let biased_cg_input_sum = self
176                .cell_gate
177                .gate_product(input_t.clone(), hidden_state.clone());
178            let candidate_cell_values = biased_cg_input_sum.tanh();
179
180            cell_state = forget_values * cell_state.clone() + add_values * candidate_cell_values;
181            hidden_state = output_values * cell_state.clone().tanh();
182
183            let unsqueezed_hidden_state = hidden_state.clone().unsqueeze_dim(1);
184
185            // store the hidden state for this timestep
186            batched_hidden_state = batched_hidden_state.slice_assign(
187                [0..batch_size, t..(t + 1), 0..self.d_hidden],
188                unsqueezed_hidden_state.clone(),
189            );
190        }
191
192        (
193            batched_hidden_state,
194            LstmState::new(cell_state, hidden_state),
195        )
196    }
197}
198
199/// Configuration to create a [BiLstm](BiLstm) module using the [init function](BiLstmConfig::init).
200#[derive(Config, Debug)]
201pub struct BiLstmConfig {
202    /// The size of the input features.
203    pub d_input: usize,
204    /// The size of the hidden state.
205    pub d_hidden: usize,
206    /// If a bias should be applied during the BiLstm transformation.
207    pub bias: bool,
208    /// BiLstm initializer
209    #[config(default = "Initializer::XavierNormal{gain:1.0}")]
210    pub initializer: Initializer,
211}
212
213/// The BiLstm module. This implementation is for Bidirectional LSTM.
214///
215/// Introduced in the paper: [Framewise phoneme classification with bidirectional LSTM and other neural network architectures](https://www.cs.toronto.edu/~graves/ijcnn_2005.pdf).
216///
217/// Should be created with [BiLstmConfig].
218#[derive(Module, Debug)]
219#[module(custom_display)]
220pub struct BiLstm<B: Backend> {
221    /// LSTM for the forward direction.
222    pub forward: Lstm<B>,
223    /// LSTM for the reverse direction.
224    pub reverse: Lstm<B>,
225    /// The size of the hidden state.
226    pub d_hidden: usize,
227}
228
229impl<B: Backend> ModuleDisplay for BiLstm<B> {
230    fn custom_settings(&self) -> Option<DisplaySettings> {
231        DisplaySettings::new()
232            .with_new_line_after_attribute(false)
233            .optional()
234    }
235
236    fn custom_content(&self, content: Content) -> Option<Content> {
237        let [d_input, _] = self
238            .forward
239            .input_gate
240            .input_transform
241            .weight
242            .shape()
243            .dims();
244        let bias = self.forward.input_gate.input_transform.bias.is_some();
245
246        content
247            .add("d_input", &d_input)
248            .add("d_hidden", &self.d_hidden)
249            .add("bias", &bias)
250            .optional()
251    }
252}
253
254impl BiLstmConfig {
255    /// Initialize a new [Bidirectional LSTM](BiLstm) module.
256    pub fn init<B: Backend>(&self, device: &B::Device) -> BiLstm<B> {
257        BiLstm {
258            forward: LstmConfig::new(self.d_input, self.d_hidden, self.bias)
259                .with_initializer(self.initializer.clone())
260                .init(device),
261            reverse: LstmConfig::new(self.d_input, self.d_hidden, self.bias)
262                .with_initializer(self.initializer.clone())
263                .init(device),
264            d_hidden: self.d_hidden,
265        }
266    }
267}
268
269impl<B: Backend> BiLstm<B> {
270    /// Applies the forward pass on the input tensor. This Bidirectional LSTM implementation
271    /// returns the state for each element in a sequence (i.e., across seq_length) and a final state.
272    ///
273    /// ## Parameters:
274    /// - batched_input: The input tensor of shape `[batch_size, sequence_length, input_size]`.
275    /// - state: An optional `LstmState` representing the initial cell state and hidden state.
276    ///   Each state tensor has shape `[2, batch_size, hidden_size]`.
277    ///   If no initial state is provided, these tensors are initialized to zeros.
278    ///
279    /// ## Returns:
280    /// - output: A tensor represents the output features of LSTM. Shape: `[batch_size, sequence_length, hidden_size * 2]`
281    /// - state: A `LstmState` represents the final forward and reverse states. Both `state.cell` and
282    ///   `state.hidden` have the shape `[2, batch_size, hidden_size]`.
283    pub fn forward(
284        &self,
285        batched_input: Tensor<B, 3>,
286        state: Option<LstmState<B, 3>>,
287    ) -> (Tensor<B, 3>, LstmState<B, 3>) {
288        let device = batched_input.clone().device();
289        let [batch_size, seq_length, _] = batched_input.shape().dims();
290
291        let [init_state_forward, init_state_reverse] = match state {
292            Some(state) => {
293                let cell_state_forward = state
294                    .cell
295                    .clone()
296                    .slice([0..1, 0..batch_size, 0..self.d_hidden])
297                    .squeeze_dim(0);
298                let hidden_state_forward = state
299                    .hidden
300                    .clone()
301                    .slice([0..1, 0..batch_size, 0..self.d_hidden])
302                    .squeeze_dim(0);
303                let cell_state_reverse = state
304                    .cell
305                    .slice([1..2, 0..batch_size, 0..self.d_hidden])
306                    .squeeze_dim(0);
307                let hidden_state_reverse = state
308                    .hidden
309                    .slice([1..2, 0..batch_size, 0..self.d_hidden])
310                    .squeeze_dim(0);
311
312                [
313                    Some(LstmState::new(cell_state_forward, hidden_state_forward)),
314                    Some(LstmState::new(cell_state_reverse, hidden_state_reverse)),
315                ]
316            }
317            None => [None, None],
318        };
319
320        // forward direction
321        let (batched_hidden_state_forward, final_state_forward) = self
322            .forward
323            .forward(batched_input.clone(), init_state_forward);
324
325        // reverse direction
326        let (batched_hidden_state_reverse, final_state_reverse) = self.reverse.forward_iter(
327            batched_input.iter_dim(1).rev().zip((0..seq_length).rev()),
328            init_state_reverse,
329            batch_size,
330            seq_length,
331            &device,
332        );
333
334        let output = Tensor::cat(
335            [batched_hidden_state_forward, batched_hidden_state_reverse].to_vec(),
336            2,
337        );
338
339        let state = LstmState::new(
340            Tensor::stack(
341                [final_state_forward.cell, final_state_reverse.cell].to_vec(),
342                0,
343            ),
344            Tensor::stack(
345                [final_state_forward.hidden, final_state_reverse.hidden].to_vec(),
346                0,
347            ),
348        );
349
350        (output, state)
351    }
352}
353
354#[cfg(test)]
355mod tests {
356    use super::*;
357    use crate::{LinearRecord, TestBackend};
358    use burn::module::Param;
359    use burn::tensor::{Device, Distribution, TensorData};
360    use burn::tensor::{ElementConversion, Tolerance, ops::FloatElem};
361    type FT = FloatElem<TestBackend>;
362
363    #[cfg(feature = "std")]
364    use crate::TestAutodiffBackend;
365
366    #[test]
367    fn test_with_uniform_initializer() {
368        let device = Default::default();
369        TestBackend::seed(&device, 0);
370
371        let config = LstmConfig::new(5, 5, false)
372            .with_initializer(Initializer::Uniform { min: 0.0, max: 1.0 });
373        let lstm = config.init::<TestBackend>(&Default::default());
374
375        let gate_to_data =
376            |gate: GateController<TestBackend>| gate.input_transform.weight.val().to_data();
377
378        gate_to_data(lstm.input_gate).assert_within_range::<FT>(0.elem()..1.elem());
379        gate_to_data(lstm.forget_gate).assert_within_range::<FT>(0.elem()..1.elem());
380        gate_to_data(lstm.output_gate).assert_within_range::<FT>(0.elem()..1.elem());
381        gate_to_data(lstm.cell_gate).assert_within_range::<FT>(0.elem()..1.elem());
382    }
383
384    /// Test forward pass with simple input vector.
385    ///
386    /// f_t = sigmoid(0.7*0.1 + 0.7*0) = sigmoid(0.07) = 0.5173928
387    /// i_t = sigmoid(0.5*0.1 + 0.5*0) = sigmoid(0.05) = 0.5123725
388    /// o_t = sigmoid(1.1*0.1 + 1.1*0) = sigmoid(0.11) = 0.5274723
389    /// c_t = tanh(0.9*0.1 + 0.9*0) = tanh(0.09) = 0.0892937
390    /// C_t = f_t * 0 + i_t * c_t = 0 + 0.5123725 * 0.0892937 = 0.04575243
391    /// h_t = o_t * tanh(C_t) = 0.5274723 * tanh(0.04575243) = 0.5274723 * 0.04568173 = 0.024083648
392    #[test]
393    fn test_forward_single_input_single_feature() {
394        let device = Default::default();
395        TestBackend::seed(&device, 0);
396
397        let config = LstmConfig::new(1, 1, false);
398        let device = Default::default();
399        let mut lstm = config.init::<TestBackend>(&device);
400
401        fn create_gate_controller(
402            weights: f32,
403            biases: f32,
404            d_input: usize,
405            d_output: usize,
406            bias: bool,
407            initializer: Initializer,
408            device: &Device<TestBackend>,
409        ) -> GateController<TestBackend> {
410            let record_1 = LinearRecord {
411                weight: Param::from_data(TensorData::from([[weights]]), device),
412                bias: Some(Param::from_data(TensorData::from([biases]), device)),
413            };
414            let record_2 = LinearRecord {
415                weight: Param::from_data(TensorData::from([[weights]]), device),
416                bias: Some(Param::from_data(TensorData::from([biases]), device)),
417            };
418            GateController::create_with_weights(
419                d_input,
420                d_output,
421                bias,
422                initializer,
423                record_1,
424                record_2,
425            )
426        }
427
428        lstm.input_gate = create_gate_controller(
429            0.5,
430            0.0,
431            1,
432            1,
433            false,
434            Initializer::XavierUniform { gain: 1.0 },
435            &device,
436        );
437        lstm.forget_gate = create_gate_controller(
438            0.7,
439            0.0,
440            1,
441            1,
442            false,
443            Initializer::XavierUniform { gain: 1.0 },
444            &device,
445        );
446        lstm.cell_gate = create_gate_controller(
447            0.9,
448            0.0,
449            1,
450            1,
451            false,
452            Initializer::XavierUniform { gain: 1.0 },
453            &device,
454        );
455        lstm.output_gate = create_gate_controller(
456            1.1,
457            0.0,
458            1,
459            1,
460            false,
461            Initializer::XavierUniform { gain: 1.0 },
462            &device,
463        );
464
465        // single timestep with single feature
466        let input = Tensor::<TestBackend, 3>::from_data(TensorData::from([[[0.1]]]), &device);
467
468        let (output, state) = lstm.forward(input, None);
469
470        let expected = TensorData::from([[0.046]]);
471        let tolerance = Tolerance::default();
472        state
473            .cell
474            .to_data()
475            .assert_approx_eq::<FT>(&expected, tolerance);
476
477        let expected = TensorData::from([[0.0242]]);
478        state
479            .hidden
480            .to_data()
481            .assert_approx_eq::<FT>(&expected, tolerance);
482
483        output
484            .select(0, Tensor::arange(0..1, &device))
485            .squeeze_dim::<2>(0)
486            .to_data()
487            .assert_approx_eq::<FT>(&state.hidden.to_data(), tolerance);
488    }
489
490    #[test]
491    fn test_batched_forward_pass() {
492        let device = Default::default();
493        let lstm = LstmConfig::new(64, 1024, true).init(&device);
494        let batched_input =
495            Tensor::<TestBackend, 3>::random([8, 10, 64], Distribution::Default, &device);
496
497        let (output, state) = lstm.forward(batched_input, None);
498
499        assert_eq!(output.dims(), [8, 10, 1024]);
500        assert_eq!(state.cell.dims(), [8, 1024]);
501        assert_eq!(state.hidden.dims(), [8, 1024]);
502    }
503
504    #[test]
505    fn test_batched_forward_pass_batch_of_one() {
506        let device = Default::default();
507        let lstm = LstmConfig::new(64, 1024, true).init(&device);
508        let batched_input =
509            Tensor::<TestBackend, 3>::random([1, 2, 64], Distribution::Default, &device);
510
511        let (output, state) = lstm.forward(batched_input, None);
512
513        assert_eq!(output.dims(), [1, 2, 1024]);
514        assert_eq!(state.cell.dims(), [1, 1024]);
515        assert_eq!(state.hidden.dims(), [1, 1024]);
516    }
517
518    #[test]
519    #[cfg(feature = "std")]
520    fn test_batched_backward_pass() {
521        use burn::tensor::Shape;
522        let device = Default::default();
523        let lstm = LstmConfig::new(64, 32, true).init(&device);
524        let shape: Shape = [8, 10, 64].into();
525        let batched_input =
526            Tensor::<TestAutodiffBackend, 3>::random(shape, Distribution::Default, &device);
527
528        let (output, _) = lstm.forward(batched_input.clone(), None);
529        let fake_loss = output;
530        let grads = fake_loss.backward();
531
532        let some_gradient = lstm
533            .output_gate
534            .hidden_transform
535            .weight
536            .grad(&grads)
537            .unwrap();
538
539        // Asserts that the gradients exist and are non-zero
540        assert_ne!(
541            some_gradient
542                .any()
543                .into_data()
544                .iter::<f32>()
545                .next()
546                .unwrap(),
547            0.0
548        );
549    }
550
551    #[test]
552    fn test_bidirectional() {
553        let device = Default::default();
554        TestBackend::seed(&device, 0);
555
556        let config = BiLstmConfig::new(2, 3, true);
557        let device = Default::default();
558        let mut lstm = config.init(&device);
559
560        fn create_gate_controller<const D1: usize, const D2: usize>(
561            input_weights: [[f32; D1]; D2],
562            input_biases: [f32; D1],
563            hidden_weights: [[f32; D1]; D1],
564            hidden_biases: [f32; D1],
565            device: &Device<TestBackend>,
566        ) -> GateController<TestBackend> {
567            let d_input = input_weights[0].len();
568            let d_output = input_weights.len();
569
570            let input_record = LinearRecord {
571                weight: Param::from_data(TensorData::from(input_weights), device),
572                bias: Some(Param::from_data(TensorData::from(input_biases), device)),
573            };
574            let hidden_record = LinearRecord {
575                weight: Param::from_data(TensorData::from(hidden_weights), device),
576                bias: Some(Param::from_data(TensorData::from(hidden_biases), device)),
577            };
578            GateController::create_with_weights(
579                d_input,
580                d_output,
581                true,
582                Initializer::XavierUniform { gain: 1.0 },
583                input_record,
584                hidden_record,
585            )
586        }
587
588        let input = Tensor::<TestBackend, 3>::from_data(
589            TensorData::from([[
590                [0.949, -0.861],
591                [0.892, 0.927],
592                [-0.173, -0.301],
593                [-0.081, 0.992],
594            ]]),
595            &device,
596        );
597        let h0 = Tensor::<TestBackend, 3>::from_data(
598            TensorData::from([[[0.280, 0.360, -1.242]], [[-0.588, 0.729, -0.788]]]),
599            &device,
600        );
601        let c0 = Tensor::<TestBackend, 3>::from_data(
602            TensorData::from([[[0.723, 0.397, -0.262]], [[0.471, 0.613, 1.885]]]),
603            &device,
604        );
605
606        lstm.forward.input_gate = create_gate_controller(
607            [[0.367, 0.091, 0.342], [0.322, 0.533, 0.059]],
608            [-0.196, 0.354, 0.209],
609            [
610                [-0.320, 0.232, -0.165],
611                [0.093, -0.572, -0.315],
612                [-0.467, 0.325, 0.046],
613            ],
614            [0.181, -0.190, -0.245],
615            &device,
616        );
617
618        lstm.forward.forget_gate = create_gate_controller(
619            [[-0.342, -0.084, -0.420], [-0.432, 0.119, 0.191]],
620            [0.315, -0.413, -0.041],
621            [
622                [0.453, 0.063, 0.561],
623                [0.211, 0.149, 0.213],
624                [-0.499, -0.158, 0.068],
625            ],
626            [-0.431, -0.535, 0.125],
627            &device,
628        );
629
630        lstm.forward.cell_gate = create_gate_controller(
631            [[-0.046, -0.382, 0.321], [-0.533, 0.558, 0.004]],
632            [-0.358, 0.282, -0.078],
633            [
634                [-0.358, 0.109, 0.139],
635                [-0.345, 0.091, -0.368],
636                [-0.508, 0.221, -0.507],
637            ],
638            [0.502, -0.509, -0.247],
639            &device,
640        );
641
642        lstm.forward.output_gate = create_gate_controller(
643            [[-0.577, -0.359, 0.216], [-0.550, 0.268, 0.243]],
644            [-0.227, -0.274, 0.039],
645            [
646                [-0.383, 0.449, 0.222],
647                [-0.357, -0.093, 0.449],
648                [-0.106, 0.236, 0.360],
649            ],
650            [-0.361, -0.209, -0.454],
651            &device,
652        );
653
654        lstm.reverse.input_gate = create_gate_controller(
655            [[-0.055, 0.506, 0.247], [-0.369, 0.178, -0.258]],
656            [0.540, -0.164, 0.033],
657            [
658                [0.159, 0.180, -0.037],
659                [-0.443, 0.485, -0.488],
660                [0.098, -0.085, -0.140],
661            ],
662            [-0.510, 0.105, 0.114],
663            &device,
664        );
665
666        lstm.reverse.forget_gate = create_gate_controller(
667            [[-0.154, -0.432, -0.547], [-0.369, -0.310, -0.175]],
668            [0.141, 0.004, 0.055],
669            [
670                [-0.005, -0.277, -0.515],
671                [-0.011, -0.101, -0.365],
672                [0.426, 0.379, 0.337],
673            ],
674            [-0.382, 0.331, -0.176],
675            &device,
676        );
677
678        lstm.reverse.cell_gate = create_gate_controller(
679            [[-0.571, 0.228, -0.287], [-0.331, 0.110, 0.219]],
680            [-0.206, -0.546, 0.462],
681            [
682                [0.449, -0.240, 0.071],
683                [-0.045, 0.131, 0.124],
684                [0.138, -0.201, 0.191],
685            ],
686            [-0.030, 0.211, -0.352],
687            &device,
688        );
689
690        lstm.reverse.output_gate = create_gate_controller(
691            [[0.491, -0.442, 0.333], [0.313, -0.121, -0.070]],
692            [-0.387, -0.250, 0.066],
693            [
694                [-0.030, 0.268, 0.299],
695                [-0.019, -0.280, -0.314],
696                [0.466, -0.365, -0.248],
697            ],
698            [-0.398, -0.199, -0.566],
699            &device,
700        );
701
702        let expected_output_with_init_state = TensorData::from([[
703            [0.23764, -0.03442, 0.04414, -0.15635, -0.03366, -0.05798],
704            [0.00473, -0.02254, 0.02988, -0.16510, -0.00306, 0.08742],
705            [0.06210, -0.06509, -0.05339, -0.01710, 0.02091, 0.16012],
706            [-0.03420, 0.07774, -0.09774, -0.02604, 0.12584, 0.20872],
707        ]]);
708        let expected_output_without_init_state = TensorData::from([[
709            [0.08679, -0.08776, -0.00528, -0.15969, -0.05322, -0.08863],
710            [-0.02577, -0.05057, 0.00033, -0.17558, -0.03679, 0.03142],
711            [0.02942, -0.07411, -0.06044, -0.03601, -0.09998, 0.04846],
712            [-0.04026, 0.07178, -0.10189, -0.07349, -0.04576, 0.05550],
713        ]]);
714        let expected_hn_with_init_state = TensorData::from([
715            [[-0.03420, 0.07774, -0.09774]],
716            [[-0.15635, -0.03366, -0.05798]],
717        ]);
718        let expected_cn_with_init_state = TensorData::from([
719            [[-0.13593, 0.17125, -0.22395]],
720            [[-0.45425, -0.11206, -0.12908]],
721        ]);
722        let expected_hn_without_init_state = TensorData::from([
723            [[-0.04026, 0.07178, -0.10189]],
724            [[-0.15969, -0.05322, -0.08863]],
725        ]);
726        let expected_cn_without_init_state = TensorData::from([
727            [[-0.15839, 0.15923, -0.23569]],
728            [[-0.47407, -0.17493, -0.19643]],
729        ]);
730
731        let (output_with_init_state, state_with_init_state) =
732            lstm.forward(input.clone(), Some(LstmState::new(c0, h0)));
733        let (output_without_init_state, state_without_init_state) = lstm.forward(input, None);
734
735        let tolerance = Tolerance::permissive();
736        output_with_init_state
737            .to_data()
738            .assert_approx_eq::<FT>(&expected_output_with_init_state, tolerance);
739        output_without_init_state
740            .to_data()
741            .assert_approx_eq::<FT>(&expected_output_without_init_state, tolerance);
742        state_with_init_state
743            .hidden
744            .to_data()
745            .assert_approx_eq::<FT>(&expected_hn_with_init_state, tolerance);
746        state_with_init_state
747            .cell
748            .to_data()
749            .assert_approx_eq::<FT>(&expected_cn_with_init_state, tolerance);
750        state_without_init_state
751            .hidden
752            .to_data()
753            .assert_approx_eq::<FT>(&expected_hn_without_init_state, tolerance);
754        state_without_init_state
755            .cell
756            .to_data()
757            .assert_approx_eq::<FT>(&expected_cn_without_init_state, tolerance);
758    }
759
760    #[test]
761    fn display_lstm() {
762        let config = LstmConfig::new(2, 3, true);
763
764        let layer = config.init::<TestBackend>(&Default::default());
765
766        assert_eq!(
767            alloc::format!("{layer}"),
768            "Lstm {d_input: 2, d_hidden: 3, bias: true, params: 84}"
769        );
770    }
771
772    #[test]
773    fn display_bilstm() {
774        let config = BiLstmConfig::new(2, 3, true);
775
776        let layer = config.init::<TestBackend>(&Default::default());
777
778        assert_eq!(
779            alloc::format!("{layer}"),
780            "BiLstm {d_input: 2, d_hidden: 3, bias: true, params: 168}"
781        );
782    }
783}
burn_nn/modules/rnn/lstm.rs

burn_nn/modules/rnn/
lstm.rs
