def topk_sum(lprobs_list, k):
    """
    lprobs_list = [lprobs_1,...,lprobs_n], where:
        lprobs_1 : (batch_size x beam_size x vocab_1)
        ...
        lprobs_n : (batch_size x beam_size x vocab_n)

    Return:
        - topk_values : (batch_size x k)
            values of the topk sum of the form :
                lprobs_1[bsz, beam_idx, vocab_1_idx] + ... + lprobs_n[bsz, beam_idx, vocab_n_idx]
        - topk_idxs : (batch_size x k x n+1)
            each (n+1)-tensor being [beam_idx, vocab_1_idx, ..., vocab_n_idx]
    """
    # Reduce all lprobs to k candidates first to reduce later complexity
    # We may assume that k << vocab
    lprobs_topk_list = []
    lprobs_topk_indices_list = []
    for lprobs in lprobs_list:
        k_i = min(k, lprobs.size(-1))
        topk_values, topk_indices = torch.topk(lprobs, k=k_i)
        # topk_values : (batch_size x beam_size x k_i)
        # topk_indices : (batch_size x beam_size x k_i)
        lprobs_topk_list.append(topk_values)
        lprobs_topk_indices_list.append(topk_indices)

    # Compute all possible sums
    sum_lprobs_topk = lprobs_topk_list[0]
    for i in range(1, len(lprobs_topk_list)):
        unsqueezed_lprobs = lprobs_topk_list[i]
        for _ in range(i):
            unsqueezed_lprobs = unsqueezed_lprobs.unsqueeze(-2)
        sum_lprobs_topk = sum_lprobs_topk.unsqueeze(-1) + unsqueezed_lprobs
    # sum_lprobs : (batch_size x beam_size x k_1 x ... x k_n)

    # Get the top k sums and the (transformed indices)
    topk_sum_values, topk_sum_indices = torch.topk(
        sum_lprobs_topk.view(sum_lprobs_topk.size(0), -1), k=k
    )
    # topk_sum_values : (batch_size x k)
    # topk_sum_indices : (batch_size x k)
    topk_sum_indices = unravel_index(topk_sum_indices, tuple(sum_lprobs_topk.shape[1:]))
    # topk_sum_indices : (batch_size x k x n+1)

    # Convert the transformed indices to the true indices
    for i_batch in range(topk_sum_indices.size(0)):
        for i_cand in range(topk_sum_indices.size(1)):
            i_beam, *transformed_vocab_indices = topk_sum_indices[i_batch, i_cand]
            true_vocab_indices = [i_beam]
            for j, transformed_vocab_j_idx in enumerate(transformed_vocab_indices):
                true_vocab_j_idx = lprobs_topk_indices_list[j][
                    i_batch, i_beam, transformed_vocab_j_idx
                ]
                true_vocab_indices.append(true_vocab_j_idx)
            topk_sum_indices[i_batch, i_cand] = torch.tensor(true_vocab_indices)

    topk_sum_beams = topk_sum_indices[:, :, 0]
    topk_sum_indices = topk_sum_indices[:, :, 1:]

    return topk_sum_values, topk_sum_indices, topk_sum_beams
lprobs_list = [lprobs_1,...,lprobs_n], where:
lprobs_1 : (batch_size x beam_size x vocab_1)
...
lprobs_n : (batch_size x beam_size x vocab_n)

Return:
- topk_values : (batch_size x k)
values of the topk sum of the form :
lprobs_1[bsz, beam_idx, vocab_1_idx] + ... + lprobs_n[bsz, beam_idx, vocab_n_idx]
- topk_idxs : (batch_size x k x n+1)
each (n+1)-tensor being [beam_idx, vocab_1_idx, ..., vocab_n_idx]
fairseq/models/speech_dlm/sequence_generator/multichannel_search.py