docstr: Tests if the embedder returns the correct shape.
docstr_tok: ["Tests","if","the","embedder","returns","the","correct","shape","."]
code: def test_embedder_returns_correct_output_shape(self):
      """Tests if the embedder returns the correct shape."""
      vocab_size = 5
      embedding_size = 3
      model = models.Embedder(
          vocab_size=vocab_size, embedding_size=embedding_size
      )
      rng = jax.random.key(0)
      token_ids = np.array([[2, 4, 3], [2, 6, 3]], dtype=np.int32)
      output, _ = model.init_with_output(
          rng, token_ids, deterministic=True
      )
      self.assertEqual(
          (token_ids.shape) + (embedding_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 1 for calling the ModelTest obj's underlying member method code with example usage: obj.test_embedder_returns_correct_output_shape() without return types