import tensorflow as tf

    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.
      tokens = tf.concat(
          [tokens[:best_id], [candidates[best_id]], tokens[best_id + 2:]],
          axis=0)
      # Recompute the merge candidates.
      # Only the neighbors of best_id need to be recomputed.
      empty = tf.zeros([0], dtype=candidates.dtype)

      def _MergeLeft():
        return tf.concat(
            [candidates[:best_id - 1],
             _MergeOneToken(tokens, best_id - 1)],
            axis=0)

      left_candidates = tf.cond(tf.equal(best_id, 0), lambda: empty, _MergeLeft)

      def _MergeRight():
        return tf.concat(
            [_MergeOneToken(tokens, best_id), candidates[best_id + 2:]], axis=0)
