docstr: Tests if the simple LSTM returns the correct shape.
docstr_tok: ["Tests","if","the","simple","LSTM","returns","the","correct","shape","."]
code: def test_lstm_returns_correct_output_shape(self):
      """Tests if the simple LSTM returns the correct shape."""
      batch_size = 2
      seq_len = 3
      embedding_size = 4
      hidden_size = 5
      model = models.SimpleLSTM(5)
      rng = jax.random.key(0)
      inputs = np.random.RandomState(0).normal(
          size=[batch_size, seq_len, embedding_size]
      )
      initial_state = model.initialize_carry(inputs[:, 0].shape)
      (_, output), _ = model.init_with_output(
          rng, initial_state, inputs
      )
      self.assertEqual((batch_size, seq_len, hidden_size), output.shape)
  
file_path: flax/examples/sst2/models_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import models
  import jax
  import jax.test_util
  import numpy
label_desc: Use image node_id 2 for calling the ModelTest obj's underlying member method code with example usage: obj.test_lstm_returns_correct_output_shape() without return types