        def noisy_channel_rescoring(lprobs, beam_size, bsz, src_tokens, tokens, k):
            """Rescore the top k hypothesis from each beam using noisy channel modeling
            Returns:
                new_fw_lprobs: the direct model probabilities after pruning the top k
                new_ch_lm_lprobs:  the combined channel and language model probabilities
                new_lm_lprobs: the language model probabilities after pruning the top k
            """
            with torch.no_grad():
                lprobs_size = lprobs.size()
                if prefix_tokens is not None and step < prefix_tokens.size(1):
                    probs_slice = lprobs.view(bsz, -1, lprobs.size(-1))[:, 0, :]
                    cand_scores = torch.gather(
                        probs_slice, dim=1,
                        index=prefix_tokens[:, step].view(-1, 1).data
                    ).expand(-1, beam_size).contiguous().view(bsz*beam_size, 1)
                    cand_indices = prefix_tokens[:, step].view(-1, 1).expand(bsz, beam_size).data.contiguous().view(bsz*beam_size, 1)

                    # need to calculate and save fw and lm probs for prefix tokens
                    fw_top_k = cand_scores
                    fw_top_k_idx = cand_indices
                    k = 1
                else:
                    # take the top k best words for every sentence in batch*beam
                    fw_top_k, fw_top_k_idx = torch.topk(lprobs.view(beam_size*bsz, -1), k=k)
                eos_idx = torch.nonzero(fw_top_k_idx.view(bsz*beam_size*k, -1) == self.eos)[:, 0]
                ch_scores = fw_top_k.new_full((beam_size*bsz*k, ), 0)
                src_size = torch.sum(src_tokens[:, :] != self.src_dict.pad_index, dim=1, keepdim=True, dtype=fw_top_k.dtype)

                if self.combine_method != "lm_only":
                    temp_src_tokens_full = src_tokens[:, :].repeat(1, k).view(bsz*beam_size*k, -1)
                    not_padding = temp_src_tokens_full[:, 1:] != self.src_dict.pad_index
                    cur_tgt_size = step+2

                    # add eos to all candidate sentences except those that already end in eos
                    eos_tokens = tokens[:, 0].repeat(1, k).view(-1, 1)
                    eos_tokens[eos_idx] = self.tgt_dict.pad_index

                    if step == 0:
                        channel_input = torch.cat((fw_top_k_idx.view(-1, 1), eos_tokens), 1)
                    else:
                        # move eos from beginning to end of target sentence
                        channel_input = torch.cat((tokens[:, 1:step + 1].repeat(1, k).view(-1, step), fw_top_k_idx.view(-1, 1), eos_tokens), 1)

                    ch_input_lengths = torch.tensor(np.full(channel_input.size(0), cur_tgt_size))
                    ch_input_lengths[eos_idx] = cur_tgt_size-1
                    if self.channel_scoring_type == "unnormalized":
                        ch_encoder_output = channel_model.encoder(channel_input, src_lengths=ch_input_lengths)
                        ch_decoder_output, _ = channel_model.decoder(temp_src_tokens_full, encoder_out=ch_encoder_output, features_only=True)
                        del ch_encoder_output
                        ch_intermed_scores = channel_model.decoder.unnormalized_scores_given_target(ch_decoder_output, target_ids=temp_src_tokens_full[:, 1:])
                        ch_intermed_scores = ch_intermed_scores.float()
                        ch_intermed_scores *= not_padding.float()
                        ch_scores = torch.sum(ch_intermed_scores, dim=1)
                    elif self.channel_scoring_type == "k2_separate":
                        for k_idx in range(k):
                            k_eos_tokens = eos_tokens[k_idx::k, :]
                            if step == 0:
                                k_ch_input = torch.cat((fw_top_k_idx[:, k_idx:k_idx+1], k_eos_tokens), 1)
                            else:
                                # move eos from beginning to end of target sentence
                                k_ch_input = torch.cat((tokens[:, 1:step + 1], fw_top_k_idx[:, k_idx:k_idx+1], k_eos_tokens), 1)
                            k_ch_input_lengths = ch_input_lengths[k_idx::k]
                            k_ch_output = channel_model(k_ch_input, k_ch_input_lengths, src_tokens)
                            k_ch_lprobs = channel_model.get_normalized_probs(k_ch_output, log_probs=True)
                            k_ch_intermed_scores = torch.gather(k_ch_lprobs[:, :-1, :], 2, src_tokens[:, 1:].unsqueeze(2)).squeeze(2)
                            k_ch_intermed_scores *= not_padding.float()
                            ch_scores[k_idx::k] = torch.sum(k_ch_intermed_scores, dim=1)
                    elif self.channel_scoring_type == "src_vocab":
                        ch_encoder_output = channel_model.encoder(channel_input, src_lengths=ch_input_lengths)
                        ch_decoder_output, _ = channel_model.decoder(temp_src_tokens_full, encoder_out=ch_encoder_output, features_only=True)

                        del ch_encoder_output
                        ch_lprobs = normalized_scores_with_batch_vocab(
                            channel_model.decoder,
                            ch_decoder_output, src_tokens, k, bsz, beam_size,
                            self.src_dict.pad_index, top_k=self.top_k_vocab)
                        ch_scores = torch.sum(ch_lprobs, dim=1)
                    elif self.channel_scoring_type == "src_vocab_batched":
                        ch_bsz_size = temp_src_tokens_full.shape[0]
                        ch_lprobs_list = [None] * len(range(0, ch_bsz_size, self.ch_scoring_bsz))
                        for i, start_idx in enumerate(range(0, ch_bsz_size, self.ch_scoring_bsz)):
                            end_idx = min(start_idx + self.ch_scoring_bsz, ch_bsz_size)
                            temp_src_tokens_full_batch = temp_src_tokens_full[start_idx:end_idx, :]
                            channel_input_batch = channel_input[start_idx:end_idx, :]
                            ch_input_lengths_batch = ch_input_lengths[start_idx:end_idx]
                            ch_encoder_output_batch = channel_model.encoder(channel_input_batch, src_lengths=ch_input_lengths_batch)
                            ch_decoder_output_batch, _ = channel_model.decoder(temp_src_tokens_full_batch, encoder_out=ch_encoder_output_batch, features_only=True)
                            ch_lprobs_list[i] = normalized_scores_with_batch_vocab(
                                channel_model.decoder,
                                ch_decoder_output_batch, src_tokens, k, bsz, beam_size,
                                self.src_dict.pad_index, top_k=self.top_k_vocab,
                                start_idx=start_idx, end_idx=end_idx)
                        ch_lprobs = torch.cat(ch_lprobs_list, dim=0)
                        ch_scores = torch.sum(ch_lprobs, dim=1)
                    else:
                        ch_output = channel_model(channel_input, ch_input_lengths, temp_src_tokens_full)
                        ch_lprobs = channel_model.get_normalized_probs(ch_output, log_probs=True)
                        ch_intermed_scores = torch.gather(ch_lprobs[:, :-1, :], 2, temp_src_tokens_full[:, 1:].unsqueeze(2)).squeeze().view(bsz*beam_size*k, -1)
                        ch_intermed_scores *= not_padding.float()
                        ch_scores = torch.sum(ch_intermed_scores, dim=1)

                else:
                    cur_tgt_size = 0
                ch_scores = ch_scores.view(bsz*beam_size, k)
                expanded_lm_prefix_scores = lm_prefix_scores.unsqueeze(1).expand(-1, k).flatten()

                if self.share_tgt_dict:
                    lm_scores = get_lm_scores(lm, tokens[:, :step + 1].view(-1, step+1), lm_incremental_states, fw_top_k_idx.view(-1, 1), torch.tensor(np.full(tokens.size(0), step+1)), k)
                else:
                    new_lm_input = dict2dict(tokens[:, :step + 1].view(-1, step+1), self.tgt_to_lm)
                    new_cands = dict2dict(fw_top_k_idx.view(-1, 1), self.tgt_to_lm)
                    lm_scores = get_lm_scores(lm, new_lm_input, lm_incremental_states, new_cands, torch.tensor(np.full(tokens.size(0), step+1)), k)

                lm_scores.add_(expanded_lm_prefix_scores)
                ch_lm_scores = combine_ch_lm(self.combine_method, ch_scores, lm_scores, src_size, cur_tgt_size)
                # initialize all as min value
                new_fw_lprobs = ch_scores.new(lprobs_size).fill_(-1e17).view(bsz*beam_size, -1)
                new_ch_lm_lprobs = ch_scores.new(lprobs_size).fill_(-1e17).view(bsz*beam_size, -1)
                new_lm_lprobs = ch_scores.new(lprobs_size).fill_(-1e17).view(bsz*beam_size, -1)
                new_fw_lprobs[:, self.pad] = -math.inf
                new_ch_lm_lprobs[:, self.pad] = -math.inf
                new_lm_lprobs[:, self.pad] = -math.inf

                new_fw_lprobs.scatter_(1, fw_top_k_idx, fw_top_k)
                new_ch_lm_lprobs.scatter_(1, fw_top_k_idx, ch_lm_scores)
                new_lm_lprobs.scatter_(1, fw_top_k_idx, lm_scores.view(-1, k))
                return new_fw_lprobs, new_ch_lm_lprobs, new_lm_lprobs
Rescore the top k hypothesis from each beam using noisy channel modeling
Returns:
new_fw_lprobs: the direct model probabilities after pruning the top k
new_ch_lm_lprobs:  the combined channel and language model probabilities
new_lm_lprobs: the language model probabilities after pruning the top k
examples/fast_noisy_channel/noisy_channel_sequence_generator.py