docstr: Tests if the simple BiLSTM returns the correct shape.
docstr_tok: ["Tests","if","the","simple","BiLSTM","returns","the","correct","shape","."]
code: def test_bilstm_returns_correct_output_shape(self):
      """Tests if the simple BiLSTM returns the correct shape."""
      batch_size = 2
      seq_len = 3
      embedding_size = 4
      hidden_size = 5
      model = models.SimpleBiLSTM(hidden_size=hidden_size)
      rng = jax.random.key(0)
      inputs = np.random.RandomState(0).normal(
          size=[batch_size, seq_len, embedding_size]
      )
      lengths = np.array([2, 3], dtype=np.int32)
      outputs, _ = model.init_with_output(rng, inputs, lengths)
      # We expect 2*hidden_size because we concatenate forward+backward LSTMs.
      self.assertEqual(
          (batch_size, seq_len, 2 * hidden_size), outputs.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 3 for calling the ModelTest obj's underlying member method code with example usage: obj.test_bilstm_returns_correct_output_shape() without return types