    def _prefix_tokens(
        self, step: int, lprobs, scores, tokens, prefix_tokens, beam_size: int
    ):
        """Handle prefix tokens"""
        prefix_toks = prefix_tokens[:, step].unsqueeze(-1).repeat(1, beam_size).view(-1)
        prefix_lprobs = lprobs.gather(-1, prefix_toks.unsqueeze(-1))
        prefix_mask = prefix_toks.ne(self.pad)
        # used for 1-channel generation, do not force the unk token (i.e. unk tokens are changed)
        prefix_mask &= prefix_toks.ne(self.unk)
        # zeroing the copying tokens
        # if step > 0:
        #     copy_mask = (prefix_tokens[:, step] == prefix_tokens[:, step-1]).unsqueeze(-1).repeat(1, beam_size).view(-1)
        #     prefix_lprobs[copy_mask & prefix_mask] = 0.
        lprobs[prefix_mask] = torch.tensor(-math.inf).to(lprobs)
        lprobs[prefix_mask] = lprobs[prefix_mask].scatter(
            -1, prefix_toks[prefix_mask].unsqueeze(-1), prefix_lprobs[prefix_mask]
        )
        # shouldn't stop at unk token
        unk_mask = prefix_toks.eq(self.unk)
        if len(lprobs[unk_mask]) > 0:
            # otherwise it won't assign to lprobs,
            # see: https://discuss.pytorch.org/t/how-to-mask-and-assign-a-value-to-tensor/18437
            copy_lprobs = lprobs[unk_mask][:, :]
            copy_lprobs[:, self.eos] = -math.inf
            lprobs[unk_mask] = copy_lprobs
        # if prefix includes eos, then we should make sure tokens and
        # scores are the same across all beams
        eos_mask = prefix_toks.eq(self.eos)
        if eos_mask.any():
            # validate that the first beam matches the prefix
            first_beam = tokens[eos_mask].view(-1, beam_size, tokens.size(-1))[
                :, 0, 1 : step + 1
            ]
            eos_mask_batch_dim = eos_mask.view(-1, beam_size)[:, 0]
            target_prefix = prefix_tokens[eos_mask_batch_dim][:, :step]
            assert (first_beam == target_prefix).all()

            # copy tokens, scores and lprobs from the first beam to all beams
            tokens = self.replicate_first_beam(tokens, eos_mask_batch_dim, beam_size)
            scores = self.replicate_first_beam(scores, eos_mask_batch_dim, beam_size)
            lprobs = self.replicate_first_beam(lprobs, eos_mask_batch_dim, beam_size)
        return lprobs, tokens, scores
Handle prefix tokens
fairseq/models/speech_dlm/sequence_generator/multichannel_sequence_generator.py