code: def get_bucketed_batches(
      dataset: tf.data.Dataset,
      batch_size: int,
      bucket_size: int,
      max_length: int,
      padded_shapes: Any,
      example_size_fn: Any,
      shuffle: bool = False,
      shuffle_seed: Optional[int] = None,
      drop_remainder: bool = False,
  ) -> tf.data.Dataset:
      """Returns padded batches of shuffled examples bucketed by length.
  
      This shuffles the examples randomly each epoch. The random order is
      deterministic and controlled by the seed.
  
      Batches are padded because sentences have different lengths.
      Sentences that are shorter in a batch will get 0s added at the end, until
      all sentences in the batch have the same length.
  
      For performance, examples of similar lengths are bucketed together. However,
      the contents of the buckets and their order is random each epoch, and
      controlled by the seed.
  
      Args:
        dataset: A TF Dataset with SST examples to be shuffled and batched.
        batch_size: The size of each batch. The remainder is dropped.
        bucket_size: How many different lengths go in each bucket.
        max_length: The maximum length to provide a bucket for.
        padded_shapes: A nested structure representing the shape to which the
          respective component of each input element should be padded prior to
          batching. See `tf.data.Dataset.padded_batch` for examples.
        example_size_fn: A TF function that returns the size of an example to
          determine in which bucket it goes. E.g., the sentence length.
        shuffle: Shuffle the dataset each epoch using seed.
        shuffle_seed: The seed that determines the shuffling order, with a
          different order each epoch.
        drop_remainder: Drop the last batch if it is not of size batch_size.
  
      Returns:
        A TF Dataset containing padded bucketed batches.
      """
      if shuffle:
          assert (
              shuffle_seed is not None
          ), "When shuffling you must provide a seed."
  
      # For bucket_size 8 and max length 24, we get bucket boundaries [9, 17, 25].
      bucket_boundaries = get_bucket_boundaries(bucket_size, max_length)
      logging.info("Batching bucket boundaries: %r", bucket_boundaries)
  
      # One batch size for each bucket plus one additional one (per requirement).
      bucket_batch_sizes = [batch_size] * (len(bucket_boundaries) + 1)
      bucket_fn = tf.data.experimental.bucket_by_sequence_length(
          example_size_fn,
          bucket_boundaries,
          bucket_batch_sizes,
          padded_shapes=padded_shapes,
          pad_to_bucket_boundary=True,
          drop_remainder=drop_remainder,
      )
  
      if shuffle:
          # For shuffling we need to know how many training examples we have.
          num_examples = get_num_examples(dataset)
          num_batches = num_examples // batch_size
          return (
              dataset.shuffle(
                  num_examples,
                  seed=shuffle_seed,
                  reshuffle_each_iteration=True,
              )
              .apply(bucket_fn)
              .shuffle(
                  num_batches,
                  seed=shuffle_seed,
                  reshuffle_each_iteration=True,
              )
              .prefetch(tf.data.experimental.AUTOTUNE)
          )
      return dataset.apply(bucket_fn).prefetch(
          tf.data.experimental.AUTOTUNE
      )
  
imports: from typing import Any, Dict, Optional
  from absl import logging
  import numpy
  import tensorflow
  import tensorflow_datasets
  import tensorflow_text
  import vocabulary
label_desc: Use image node_id 3 for calling a global function with example usage: get_bucketed_batches(dataset, batch_size, bucket_size, max_length, padded_shapes, example_size_fn, shuffle, shuffle_seed, drop_remainder) and returns: dataset, dataset