    def _generate(
        self,
        sample: Dict[str, Dict[str, Tensor]],
        prefix_tokens: Optional[Dict[str, Tensor]] = None,
        constraints: Optional[Tensor] = None,
        bos_token: Optional[int] = None,
    ):
        """
        Here sample is expected to have the following form
            {
                'id': index,
                'net_input': {
                    'src_tokens': {
                        'channel1' : tensor((batch x src_length)),
                        'channel2' : tensor((batch x src_length)),
                    },
                    ...
                },
            }
        and prefix_tokens
            {
                'channel1' : tensor((batch x prefix_length)),
                'channel2' : tensor((batch x prefix_length)),
            }
        """
        if self.model.is_speech_dlm:
            incremental_states = torch.jit.annotate(
                List[Dict[str, Dict[str, Optional[Tensor]]]],
                [
                    torch.jit.annotate(
                        List[Dict[str, Dict[str, Optional[Tensor]]]],
                        [{} for _ in range(self.n_channels)],
                    )
                    for i in range(self.model.models_size)
                ],
            )
        else:
            incremental_states = torch.jit.annotate(
                List[Dict[str, Dict[str, Optional[Tensor]]]],
                [
                    torch.jit.annotate(Dict[str, Dict[str, Optional[Tensor]]], {})
                    for i in range(self.model.models_size)
                ],
            )
        net_input = sample["net_input"]
        # Convert from dict to tensor form
        # shape of src_tokens : (bsz x src_len x n_channels)
        src_tokens = torch.stack(
            [net_input["src_tokens"][channel] for channel in self.channels], dim=-1
        )
        prefix_tokens = torch.stack(
            [prefix_tokens[channel] for channel in self.channels], dim=-1
        )
        # length of the source text being the character length except EndOfSentence and pad
        src_lengths = (
            (src_tokens[..., 0].ne(self.eos) & src_tokens[..., 0].ne(self.pad))
            .long()
            .sum(dim=1)
        )

        # bsz: total number of sentences in beam
        # Note that src_tokens may have more than 2 dimensions (i.e. audio features)
        bsz, src_len = src_tokens.size()[:2]
        beam_size = self.beam_size

        if constraints is not None and not self.search.supports_constraints:
            raise NotImplementedError(
                "Target-side constraints were provided, but search method doesn't support them"
            )

        # Initialize constraints, when active
        self.search.init_constraints(constraints, beam_size)

        max_len: int = -1
        if self.match_source_len:
            max_len = src_lengths.max().item()
        else:
            max_len = min(
                int(self.max_len_a * src_len + self.max_len_b),
                # exclude the EOS marker
                self.model.max_decoder_positions() - 1,
            )
        assert (
            self.min_len <= max_len
        ), "min_len cannot be larger than max_len, please adjust these!"
        # compute the encoder output for each beam
        encoder_outs = self.model.forward_encoder(net_input)

        # placeholder of indices for bsz * beam_size to hold tokens and accumulative scores
        new_order = torch.arange(bsz).view(-1, 1).repeat(1, beam_size).view(-1)
        new_order = new_order.to(src_tokens.device).long()
        encoder_outs = self.model.reorder_encoder_out(encoder_outs, new_order)
        # ensure encoder_outs is a List.
        assert encoder_outs is not None

        # initialize buffers
        # cumulative scores of hypotheses
        scores = (
            torch.zeros(bsz * beam_size, max_len + 1, self.n_channels)
            .to(src_tokens)
            .float()
        )  # +1 for eos; pad is never chosen for scoring
        tokens = (
            torch.zeros(bsz * beam_size, max_len + 2, self.n_channels)
            .to(src_tokens)
            .long()
            .fill_(self.pad)
        )  # +2 for eos and pad
        tokens[:, 0] = self.eos if bos_token is None else bos_token
        attn: Optional[Tensor] = None

        # A list that indicates candidates that should be ignored.
        # For example, suppose we're sampling and have already finalized 2/5
        # samples. Then cands_to_ignore would mark 2 positions as being ignored,
        # so that we only finalize the remaining 3 samples.
        cands_to_ignore = (
            torch.zeros(bsz, beam_size).to(src_tokens).eq(-1)
        )  # forward and backward-compatible False mask

        # list of completed sentences
        finalized = torch.jit.annotate(
            List[List[Dict[str, Tensor]]],
            [torch.jit.annotate(List[Dict[str, Tensor]], []) for i in range(bsz)],
        )  # contains lists of dictionaries of infomation about the hypothesis being finalized at each step

        finished = [
            False for i in range(bsz)
        ]  # a boolean array indicating if the sentence at the index is finished or not
        num_remaining_sent = bsz  # number of sentences remaining

        # number of candidate hypos per step
        cand_size = 2 * beam_size  # 2 x beam size in case half are EOS

        # offset arrays for converting between different indexing schemes
        bbsz_offsets = (
            (torch.arange(0, bsz) * beam_size)
            .unsqueeze(1)
            .type_as(tokens)
            .to(src_tokens.device)
        )
        cand_offsets = torch.arange(0, cand_size).type_as(tokens).to(src_tokens.device)

        reorder_state: Optional[Tensor] = None
        batch_idxs: Optional[Tensor] = None

        original_batch_idxs: Optional[Tensor] = None
        if "id" in sample and isinstance(sample["id"], Tensor):
            original_batch_idxs = sample["id"]
        else:
            original_batch_idxs = torch.arange(0, bsz).type_as(tokens)

        if self.duration_prediction:
            dur_counter = torch.ones(bsz * beam_size, self.n_channels).to(src_tokens)
            # save the indice where the dur_counter just copied from dur_pred
            dur_counter_jump_indices = None

        for step in range(max_len + 1):  # one extra step for EOS marker
            # reorder decoder internal states based on the prev choice of beams
            if reorder_state is not None:
                if batch_idxs is not None:
                    # update beam indices to take into account removed sentences
                    corr = batch_idxs - torch.arange(batch_idxs.numel()).type_as(
                        batch_idxs
                    )
                    reorder_state.view(-1, beam_size).add_(
                        corr.unsqueeze(-1) * beam_size
                    )
                    original_batch_idxs = original_batch_idxs[batch_idxs]
                self.model.reorder_incremental_state(incremental_states, reorder_state)
                encoder_outs = self.model.reorder_encoder_out(
                    encoder_outs, reorder_state
                )

            input_tokens = {
                channel: tokens[:, : step + 1, i]
                for i, channel in enumerate(self.channels)
            }

            lprobs_dict, avg_attn_scores = self.model.forward_decoder(
                input_tokens,
                encoder_outs,
                incremental_states,
                self.temperature,
            )

            # Because the sizes of vocab is different, we cannot concat the lprobs to form a single tensor
            if not self.duration_prediction:
                lprobs_list = list(lprobs_dict.values())
            else:
                lprobs_list = [
                    net_output["pred_token"] for net_output in lprobs_dict.values()
                ]

                # non-positive predicted durations
                dur_preds = (
                    torch.stack(
                        [
                            net_output["pred_duration"]
                            for net_output in lprobs_dict.values()
                        ]
                    )
                    .squeeze(-1)
                    .T
                )
                dur_preds = dur_preds / self.duration_temperature
                dur_preds = dur_preds.round().long()
                dur_preds[dur_preds < 1] = 1

                # dur_preds & dur_counter needs to be modified when there isn't an edge
                if step > 0:
                    non_edge_indices = tokens[:, step, :] == tokens[:, step - 1, :]
                    if self.delayed_duration:
                        dur_preds[non_edge_indices] = 1
                    else:
                        if dur_counter_jump_indices is not None:
                            dur_counter[dur_counter_jump_indices & non_edge_indices] = 2

                # update dur_counter
                if step > 0:
                    if self.delayed_duration:
                        dur_counter -= (
                            (dur_counter == 1)
                            | (tokens[:, step, :] == tokens[:, step - 1, :])
                        ).int()
                        dur_counter[dur_counter < 0] = 0
                    else:
                        dur_counter -= (
                            tokens[:, step, :] == tokens[:, step - 1, :]
                        ).int()
                        dur_counter[dur_counter < 1] = 1

                # whether to copy previous token (ie. if the counter is still on)
                # and get get the new duration
                if self.delayed_duration:
                    dur_counter_jump_indices = dur_counter == 0
                    dur_counter[dur_counter_jump_indices] = dur_preds[
                        dur_counter_jump_indices
                    ]

                # whether to copy previous token in this step
                copy_prev_token = dur_counter != 1
                if self.delayed_duration is False:
                    dur_counter_jump_indices = dur_counter == 1
                    dur_counter[dur_counter_jump_indices] = dur_preds[
                        dur_counter_jump_indices
                    ]
                # else:
                # dur_counter[dur_counter==0] = dur_preds[dur_counter==0] - 1
                # copy_prev_token = (dur_counter > 0)

            if self.lm_model is not None:
                assert False, "Currently not supported in multichannelLM case"

            for i in range(self.n_channels):
                lprobs_list[i][lprobs_list[i] != lprobs_list[i]] = torch.tensor(
                    -math.inf
                ).to(lprobs_list[i])

                lprobs_list[i][:, self.pad] = -math.inf  # never select pad
                lprobs_list[i][:, self.unk] -= self.unk_penalty  # apply unk penalty

                # handle max length constraint
                if step >= max_len:
                    lprobs_list[i][:, : self.eos] = -math.inf
                    lprobs_list[i][:, self.eos + 1 :] = -math.inf
                else:
                    lprobs_list[i][
                        :, self.eos
                    ] = -math.inf  # quick fix for short generation

                # handle prefix tokens (possibly with different lengths)
                if (
                    prefix_tokens is not None
                    and step < prefix_tokens.size(1)
                    and step < max_len
                ):
                    (
                        lprobs_list[i],
                        tokens[..., i],
                        scores[..., i],
                    ) = self._prefix_tokens(
                        step,
                        lprobs_list[i],
                        scores[..., i],
                        tokens[..., i],
                        prefix_tokens[..., i],
                        beam_size,
                    )
                    if self.duration_prediction:
                        # Can copy previous token if the prefix token is padding or unk (1-channel conditionned case)
                        can_copy_mask = (
                            prefix_tokens[:, step, i].eq(self.pad)
                            | prefix_tokens[:, step, i].eq(self.unk)
                        ).repeat_interleave(beam_size)
                        copy_prev_token[:, i] &= can_copy_mask
                elif step < self.min_len:
                    # minimum length constraint (does not apply if using prefix_tokens)
                    lprobs_list[i][:, self.eos] = -math.inf

                if self.duration_prediction:
                    if step < max_len:
                        for j in range(copy_prev_token.size(0)):
                            if copy_prev_token[j, i]:
                                prev_token = tokens[j, step, i]
                                lprobs_list[i][j, :prev_token] = -math.inf
                                lprobs_list[i][j, prev_token + 1 :] = -math.inf
                                # lprobs_list[i][j, prev_token] = 0.
                                # dur_counter[j,i] -= 1
                            # else:
                            #     prev_token = tokens[j, step, i]
                            # if not (lprobs_list[i][j,:].ne(-math.inf).nonzero() == prev_token).all():
                            #     lprobs_list[i][j, prev_token] = -math.inf
                            #     dur_counter[j,i] = 0.

            # Record attention scores, only support avg_attn_scores is a Tensor
            if avg_attn_scores is not None:
                if attn is None:
                    attn = torch.empty(
                        bsz * beam_size, avg_attn_scores.size(1), max_len + 2
                    ).to(scores)
                attn[:, :, step + 1].copy_(avg_attn_scores)

            scores = scores.type_as(lprobs_list[0])
            eos_bbsz_idx = torch.empty(0).to(
                tokens
            )  # indices of hypothesis ending with eos (finished sentences)
            eos_scores = torch.empty(0).to(
                scores
            )  # scores of hypothesis ending with eos (finished sentences)

            if self.should_set_src_lengths:
                self.search.set_src_lengths(src_lengths)

            if self.repeat_ngram_blocker is not None:
                for i in range(self.n_channels):
                    lprobs_list[i] = self.repeat_ngram_blocker(
                        tokens, lprobs_list[i], bsz, beam_size, step
                    )

            # Shape: (batch, cand_size)
            cand_scores, cand_indices, cand_beams = self.search.step(
                step,
                [
                    lprobs_list[i].view(bsz, -1, self.vocab_sizes[i])
                    for i in range(self.n_channels)
                ],
                scores.view(bsz, beam_size, -1, self.n_channels)[:, :, :step, :],
                tokens[:, : step + 1],
                original_batch_idxs,
            )

            # cand_bbsz_idx contains beam indices for the top candidate
            # hypotheses, with a range of values: [0, bsz*beam_size),
            # and dimensions: [bsz, cand_size]
            cand_bbsz_idx = cand_beams.add(bbsz_offsets)

            # finalize hypotheses that end in eos
            # Shape of eos_mask: (batch size, beam size)
            eos_mask = cand_indices.eq(self.eos) & cand_scores.ne(-math.inf)
            eos_mask = torch.any(eos_mask, dim=-1, keepdim=False)
            eos_mask[:, :beam_size][cands_to_ignore] = torch.tensor(0).to(eos_mask)

            # only consider eos when it's among the top beam_size indices
            # Now we know what beam item(s) to finish
            # Shape: 1d list of absolute-numbered
            eos_bbsz_idx = torch.masked_select(
                cand_bbsz_idx[:, :beam_size], mask=eos_mask[:, :beam_size]
            )

            finalized_sents: List[int] = []
            if eos_bbsz_idx.numel() > 0:
                eos_scores = torch.stack(
                    [
                        torch.masked_select(
                            cand_scores[:, :beam_size, i], mask=eos_mask[:, :beam_size]
                        )
                        for i in range(self.n_channels)
                    ],
                    dim=-1,
                )
                finalized_sents = self.finalize_hypos(
                    step,
                    eos_bbsz_idx,
                    eos_scores,
                    tokens,
                    scores,
                    finalized,
                    finished,
                    beam_size,
                    attn,
                    src_lengths,
                    max_len,
                )
                num_remaining_sent -= len(finalized_sents)

            assert num_remaining_sent >= 0
            if num_remaining_sent == 0:
                break
            if self.search.stop_on_max_len and step >= max_len:
                break
            assert step < max_len, f"{step} < {max_len}"

            # Remove finalized sentences (ones for which {beam_size}
            # finished hypotheses have been generated) from the batch.
            if len(finalized_sents) > 0:
                new_bsz = bsz - len(finalized_sents)

                # construct batch_idxs which holds indices of batches to keep for the next pass
                batch_mask = torch.ones(
                    bsz, dtype=torch.bool, device=cand_indices.device
                )
                batch_mask[finalized_sents] = False
                # TODO replace `nonzero(as_tuple=False)` after TorchScript supports it
                batch_idxs = torch.arange(
                    bsz, device=cand_indices.device
                ).masked_select(batch_mask)

                # Choose the subset of the hypothesized constraints that will continue
                self.search.prune_sentences(batch_idxs)

                eos_mask = eos_mask[batch_idxs]
                cand_beams = cand_beams[batch_idxs]
                bbsz_offsets.resize_(new_bsz, 1)
                cand_bbsz_idx = cand_beams.add(bbsz_offsets)
                cand_scores = cand_scores[batch_idxs]
                cand_indices = cand_indices[batch_idxs]

                if prefix_tokens is not None:
                    prefix_tokens = prefix_tokens[batch_idxs]
                src_lengths = src_lengths[batch_idxs]
                cands_to_ignore = cands_to_ignore[batch_idxs]

                scores = scores.view(bsz, -1)[batch_idxs].view(
                    new_bsz * beam_size, -1, self.n_channels
                )
                tokens = tokens.view(bsz, -1)[batch_idxs].view(
                    new_bsz * beam_size, -1, self.n_channels
                )
                if self.duration_prediction:
                    dur_counter = dur_counter.view(bsz, -1)[batch_idxs].view(
                        new_bsz * beam_size, self.n_channels
                    )
                if attn is not None:
                    attn = attn.view(bsz, -1)[batch_idxs].view(
                        new_bsz * beam_size, attn.size(1), -1
                    )
                bsz = new_bsz
            else:
                batch_idxs = None

            # Set active_mask so that values > cand_size indicate eos hypos
            # and values < cand_size indicate candidate active hypos.
            # After, the min values per row are the top candidate active hypos

            # Rewrite the operator since the element wise or is not supported in torchscript.

            eos_mask[:, :beam_size] = ~((~cands_to_ignore) & (~eos_mask[:, :beam_size]))
            active_mask = torch.add(
                eos_mask.type_as(cand_offsets) * cand_size,
                cand_offsets[: eos_mask.size(1)],
            )

            # get the top beam_size active hypotheses, which are just
            # the hypos with the smallest values in active_mask.
            # {active_hypos} indicates which {beam_size} hypotheses
            # from the list of {2 * beam_size} candidates were
            # selected. Shapes: (batch size, beam size)
            new_cands_to_ignore, active_hypos = torch.topk(
                active_mask, k=beam_size, dim=1, largest=False
            )

            # update cands_to_ignore to ignore any finalized hypos.
            cands_to_ignore = new_cands_to_ignore.ge(cand_size)[:, :beam_size]
            # Make sure there is at least one active item for each sentence in the batch.
            assert (~cands_to_ignore).any(dim=1).all()

            # update cands_to_ignore to ignore any finalized hypos
            # {active_bbsz_idx} denotes which beam number is continued for each new hypothesis (a beam
            # can be selected more than once).
            active_bbsz_idx = torch.gather(cand_bbsz_idx, dim=1, index=active_hypos)
            active_bbsz_idx = active_bbsz_idx.view(-1)

            # active_scores = torch.stack([
            #     torch.gather(cand_scores[...,0], dim=1, index=active_hypos)
            #         for i in range(self.n_channels)
            #         ], dim = -1)
            # active_scores = active_scores.view(-1)

            # copy tokens and scores for active hypotheses

            # Set the tokens for each beam (can select the same row more than once)
            tokens[:, : step + 1] = torch.index_select(
                tokens[:, : step + 1], dim=0, index=active_bbsz_idx
            )
            # Select the next token for each of them
            for i in range(self.n_channels):
                tokens.view(bsz, beam_size, -1, self.n_channels)[
                    :, :, step + 1, i
                ] = torch.gather(cand_indices[..., i], dim=1, index=active_hypos)
            if step > 0:
                scores[:, :step] = torch.index_select(
                    scores[:, :step], dim=0, index=active_bbsz_idx
                )
            for i in range(self.n_channels):
                scores.view(bsz, beam_size, -1, self.n_channels)[
                    :, :, step, i
                ] = torch.gather(cand_scores[..., i], dim=1, index=active_hypos)

            if self.duration_prediction:
                dur_counter = torch.index_select(
                    dur_counter, dim=0, index=active_bbsz_idx
                )

            # Update constraints based on which candidates were selected for the next beam
            self.search.update_constraints(active_hypos)

            # copy attention for active hypotheses
            if attn is not None:
                attn[:, :, : step + 2] = torch.index_select(
                    attn[:, :, : step + 2], dim=0, index=active_bbsz_idx
                )

            # reorder incremental state in decoder
            reorder_state = active_bbsz_idx

        # sort by score descending
        for sent in range(len(finalized)):
            scores = torch.tensor(
                [float(elem["score"].item()) for elem in finalized[sent]]
            )
            _, sorted_scores_indices = torch.sort(scores, descending=True)
            finalized[sent] = [finalized[sent][ssi] for ssi in sorted_scores_indices]
            finalized[sent] = torch.jit.annotate(
                List[Dict[str, Tensor]], finalized[sent]
            )
        return finalized
Here sample is expected to have the following form
{
'id': index,
'net_input': {
'src_tokens': {
'channel1' : tensor((batch x src_length)),
'channel2' : tensor((batch x src_length)),
},
...
},
}
and prefix_tokens
{
'channel1' : tensor((batch x prefix_length)),
'channel2' : tensor((batch x prefix_length)),
}
fairseq/models/speech_dlm/sequence_generator/multichannel_sequence_generator.py