docstr: Each batch should have a length that is a multiple of bucket_size.
docstr_tok: ["Each","batch","should","have","a","length","that","is","a","multiple","of","bucket_size","."]
code: def test_bucketed_dataset(self):
      """Each batch should have a length that is a multiple of bucket_size."""
      batch_size = 2
      bucket_size = 8
      for batch in self.dataset.get_bucketed_batches(
          batch_size=batch_size,
          bucket_size=bucket_size,
          max_input_length=60,
          shuffle=False,
      ).take(3):
          # Because of bucketing, sequence length must be multiple of bucket_size.
          length = batch["token_ids"].numpy().shape[-1]
          self.assertEqual(0, length % bucket_size)
          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 4 for calling the InputPipelineTest obj's underlying member method code with example usage: obj.test_bucketed_dataset() without return types