import tensorflow as tf

        tf.broadcast_to(self.unk_id, tf.shape(tokens)),
        tokens)
    # Create initial candidate list.
    candidates = tf.map_fn(
        self._MergeTokens, (tokens[:-1], tokens[1:]), dtype=tokens.dtype)

    def _ShouldMerge(unused_tokens, candidates):
      """Merge until not possible, or we abort early according to merge_prob."""
      return tf.logical_and(
          tf.reduce_any(tf.not_equal(candidates, NO_TOKEN)),
          tf.random.uniform([]) < self._merge_prob)

    def _MergeOneToken(tokens, i):
      return tf.expand_dims(
          self._MergeTokens((tokens[i], tokens[i + 1])), axis=-1)

    def _MergeCandidates(tokens, candidates):
      """Merge in the reverse binary tree."""
      best_id = tf.argmin(candidates, output_type=tf.int32)
      # Perform the merge at position best_id.
