def pack_and_batch(dataset, batch_size, length, pack=True):
  """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].
  """
  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
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].