import tensorflow as tf

@gin.configurable(module='trax.data', denylist=['dataset', 'training'])
def pad_dataset_to_length(dataset, training, len_map=None):
  """Pad features less than specified length to specified length."""
  del training
  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
