code: def batch_dataset(
      dataset,
      batch_size,
      shapes,
      max_length,
      n_epochs=1,
      shuffle=False,
      table_batching=False,
      random_seed=42,
  ):
      if isinstance(shapes, tuple):
          shapes = (
              {**shapes[0], "length": tf.TensorShape([])},
              shapes[1],
          )
      else:
          shapes = {**shapes, "length": tf.TensorShape([])}
  
      if table_batching:
          assert isinstance(
              shapes, tuple
          ), "You cannot use table batching to predict on tables as order is not guarenteed"
  
          def batched_dataset():
              return (
                  dataset()
                  .map(add_length)
                  .shuffle(500 if shuffle else 1, seed=random_seed)
                  # When we update to tf.2.13 this will change to be a method on the dataset.
                  .apply(
                      tf.data.experimental.bucket_by_sequence_length(
                          element_length_func=(
                              lambda item, *_: tf.cast(
                                  tf.maximum(
                                      tf.reduce_sum(
                                          item["context"][:, 0]
                                          - item["context"][:, 2]
                                          + 1
                                      ),
                                      tf.reduce_sum(
                                          item["context"][:, 1]
                                          - item["context"][:, 3]
                                          + 1
                                      ),
                                  ),
                                  tf.int32,
                              )
                          ),
                          bucket_boundaries=[max_length],
                          bucket_batch_sizes=[batch_size, 1],
                          padded_shapes=shapes,
                          drop_remainder=False,
                      )
                  )
                  .repeat(n_epochs)
                  .prefetch(tf.data.experimental.AUTOTUNE)
              )
  
      else:
  
          def batched_dataset():
              return (
                  dataset()
                  .map(add_length)
                  .shuffle(500 if shuffle else 1, seed=random_seed)
                  .padded_batch(
                      batch_size,
                      padded_shapes=shapes,
                      drop_remainder=False,
                  )
                  .repeat(n_epochs)
                  .prefetch(tf.data.experimental.AUTOTUNE)
              )
  
      return batched_dataset
  
file_path: finetune/finetune/util/input_utils.py
imports: import math
  import tensorflow
  from finetune.util.timing import ProgressBar
label_desc: Use image node_id 5 for calling a global function with example usage: batch_dataset(dataset, batch_size, shapes, max_length, n_epochs, shuffle, table_batching, random_seed) and returns: batched_dataset, dataset, dataset