ee44690b0c5c6757eb7eb659fdfff54c4e982da9,mesh_tensorflow/transformer/dataset.py,,pack_and_batch,#Any#Any#Any#Any#,99

Before Change

  )
  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

After Change

  if pack:
    dataset = pack_dataset(dataset, length=length)
  // Pad/trim length of each example to length.
  dataset = trim_and_pad_dataset(dataset, length=length),
  dataset = dataset.batch(batch_size, drop_remainder=False)
  // Pad final batch.
  dataset = trim_and_pad_dataset(dataset, length=batch_size),
  dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
  return dataset

Instances