code: def pack_and_batch(dataset, batch_size, length, pack=True):
    
    if pack:
      dataset = pack_dataset(dataset, length=length)
    # Pad/trim length of each example to length
    dataset = dataset.map(
        functools.partial(trim_and_pad_all_features, length=length),
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    )
    dataset = dataset.batch(batch_size, drop_remainder=False)
    # Pad batch size of each batch to batch_size
    dataset = dataset.map(
        functools.partial(trim_and_pad_all_features, length=batch_size),
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    )
    # Remind TensorFlow of the shape
    dataset = dataset.map(
        lambda x: {k: tf.reshape(v, (batch_size, length)) for k, v in x.items()},
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    )
    dataset = dataset.prefetch(100)
    return dataset
docs: Create a tf.data.Dataset which emits training batches.
  
    The input dataset emits feature-dictionaries where each feature is a vector
    of integers ending in EOS=1
  
    The tensors in the returned tf.data.Dataset have shape
    [batch_size, length].  Zeros indicate padding.
  
    length indicates the length of the emitted examples.  Examples with
    inputs/targets longer than length get truncated.
  
    TODO(noam): for text2self problems, we should just chop too-long
    sequences into multiple parts and train on all of them.
  
    If pack=False, then each emitted example will contain one
    example emitted by load_internal().
  
    If pack=True, then multiple examples emitted by load_internal() are
    concatenated to form one combined example with the given length.
    See comments in the function pack_dataset().
  
    batch_size indicates the number of (combined) examples per batch,
    across all cores.
  
    Args:
      dataset: a tf.data.Dataset
      batch_size: an integer
      length: an integer
      pack: a boolean
    Returns:
      a tf.data.Dataset where all features have fixed shape [batch, length].