def extract_best_beam_single(sample_words):
  """Extract the best beam from the sampled words."""
  if sample_words.ndim == 3:  # if this is beam search
    # Original beam search output is in [batch_size, time, beam_width] shape.
    sample_words = sample_words.transpose([2, 0, 1])
    # After extraction it would be [batch_size, time]
    best_beam = sample_words[0, :, :]
    return best_beam
  else:
    return sample_words