    def step(
        self, step, lprobs, scores, prev_output_tokens=None, original_batch_idxs=None
    ):
        """Take a single search step.

        Args:
            step: the current search step, starting at 0
            lprobs: dictionary of channels {channel : (bsz x input_beam_size x vocab_size_channel)}
                the model's log-probabilities over the vocabulary at the current step
            scores: {channel : (bsz x input_beam_size x step)}
                the historical model scores of each hypothesis up to this point
            prev_output_tokens: {channel : (bsz x step)}
                the previously generated oputput tokens
            original_batch_idxs: (bsz)
                the tensor with the batch indices, in the range [0, bsz)
                this is useful in case there has been applied a re-ordering
                and we need to know the orignal indices

        Return: A tuple of (scores, indices, beams) where:
            scores: {channel : (bsz x output_beam_size)}
                the scores of the chosen elements; output_beam_size can be
                larger than input_beam_size, e.g., we may return
                2*input_beam_size to account for EOS
            indices: {channel : (bsz x output_beam_size)}
                the indices of the chosen elements
            beams: (bsz x output_beam_size)
                the hypothesis ids of the chosen elements, in the range [0, input_beam_size)
        """
        raise NotImplementedError
Take a single search step.

Args:
step: the current search step, starting at 0
lprobs: dictionary of channels {channel : (bsz x input_beam_size x vocab_size_channel)}
the model's log-probabilities over the vocabulary at the current step
scores: {channel : (bsz x input_beam_size x step)}
the historical model scores of each hypothesis up to this point
prev_output_tokens: {channel : (bsz x step)}
the previously generated oputput tokens
original_batch_idxs: (bsz)
the tensor with the batch indices, in the range [0, bsz)
this is useful in case there has been applied a re-ordering
and we need to know the orignal indices

Return: A tuple of (scores, indices, beams) where:
scores: {channel : (bsz x output_beam_size)}
the scores of the chosen elements; output_beam_size can be
larger than input_beam_size, e.g., we may return
2*input_beam_size to account for EOS
indices: {channel : (bsz x output_beam_size)}
the indices of the chosen elements
beams: (bsz x output_beam_size)
the hypothesis ids of the chosen elements, in the range [0, input_beam_size)
fairseq/models/speech_dlm/sequence_generator/multichannel_search.py