import tensorflow as tf

  if len_map is None:
    return dataset

  def pad_to_len(x):
    for key, max_len in len_map.items():
      x_shape = tf.shape(x[key])
      x_len = x_shape[0]
      if x_len < max_len:
        pad_shape = [
            max_len - x_len,
        ]
        zeros = tf.zeros(pad_shape, dtype=x[key].dtype)
        x[key] = tf.concat([x[key], zeros], 0)
    return x

  return dataset.map(pad_to_len)


@gin.configurable(module='trax.data', denylist=['dataset', 'training'])
def add_eos_to_output_features(dataset,
                               training,
                               output_features='targets',
