docstr: Tests that the length of a batch matches the longest sequence.
docstr_tok: ["Tests","that","the","length","of","a","batch","matches","the","longest","sequence","."]
code: def test_batched_dataset(self):
      """Tests that the length of a batch matches the longest sequence."""
      batch_size = 2
      for batch in self.dataset.get_batches(
          batch_size=batch_size, shuffle=False
      ).take(1):
          # Each batch is padded to the maximum sentence length in the batch.
          max_length_in_batch = max(batch["length"].numpy())
          length = batch["token_ids"].numpy().shape[-1]
          self.assertEqual(max_length_in_batch, length)
          self.assertEqual((batch_size,), batch["length"].numpy().shape)
          self.assertEqual((batch_size,), batch["label"].numpy().shape)
  
file_path: flax/examples/sst2/input_pipeline_test.py
imports: import os
  import pathlib
  import tempfile
  from absl.testing import absltest
  import tensorflow_datasets
  import input_pipeline
  import vocabulary
label_desc: Use image node_id 5 for calling the InputPipelineTest obj's underlying member method code with example usage: obj.test_batched_dataset() without return types