def get_token_mask(token_ids, stop_id):
  """Create mask for all ids past stop_id (inclusive)."""
  batch_size = tensor_utils.shape(token_ids, 0)
  num_tokens = tensor_utils.shape(token_ids, 1)

  # Create position matrix.
  idx_range = tf.expand_dims(tf.range(num_tokens), 0)
  idx_range = tf.tile(idx_range, [batch_size, 1])

  # Find positions of stop_id.
  stop_positions = tf.where(
      condition=tf.equal(token_ids, stop_id),
      x=idx_range,
      y=tf.fill([batch_size, num_tokens], num_tokens))

  # Find earliest stop position (length).
  stop_positions = tf.reduce_min(stop_positions, -1)

  # Mask out all tokens at positions > stop_id.
  mask = tf.less_equal(idx_range, tf.expand_dims(stop_positions, -1))

  return tf.cast(mask, tf.int32)
Create mask for all ids past stop_id (inclusive).
language/capwap/utils/text_utils.py