docstr: Tests if a TextClassifier model returns the correct shape.
docstr_tok: ["Tests","if","a","TextClassifier","model","returns","the","correct","shape","."]
code: def test_text_classifier_returns_correct_output_shape(self):
      """Tests if a TextClassifier model returns the correct shape."""
      embedding_size = 3
      hidden_size = 7
      vocab_size = 5
      output_size = 3
      dropout_rate = 0.1
      word_dropout_rate = 0.2
      unk_idx = 1
  
      model = models.TextClassifier(
          embedding_size=embedding_size,
          hidden_size=hidden_size,
          vocab_size=vocab_size,
          output_size=output_size,
          dropout_rate=dropout_rate,
          word_dropout_rate=word_dropout_rate,
          unk_idx=unk_idx,
          deterministic=True,
      )
  
      rng = jax.random.key(0)
      token_ids = np.array([[2, 4, 3], [2, 6, 3]], dtype=np.int32)
      lengths = np.array([2, 3], dtype=np.int32)
      output, _ = model.init_with_output(rng, token_ids, lengths)
      batch_size = token_ids.shape[0]
      self.assertEqual((batch_size, output_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 4 for calling the ModelTest obj's underlying member method code with example usage: obj.test_text_classifier_returns_correct_output_shape() without return types