def dynamic_padding_1d(tensor: tf.Tensor,
                       length: int,
                       padding_token_id: int = 0) -> tf.Tensor:
  """Padds or truncates 1D tensor to a specified length."""
  length_to_pad = length - tf.shape(tensor)[0]
  length_to_pad_or_zero = tf.maximum(length_to_pad, 0)
  length_to_pad_or_zero = tf.expand_dims(length_to_pad_or_zero, 0)
  paddings = tf.concat([tf.constant([0]), length_to_pad_or_zero], axis=0)
  paddings = tf.expand_dims(paddings, 0)

  def pad():
    return tf.pad(
        tensor, paddings, 'CONSTANT', constant_values=padding_token_id)

  padded_tensor = tf.cond(
      length_to_pad > 0, true_fn=pad, false_fn=lambda: tensor[:length])
  padded_tensor.set_shape(length)
  return padded_tensor
Padds or truncates 1D tensor to a specified length.
language/mentionmemory/utils/mention_preprocess_utils.py