code: 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
  
imports: from typing import Dict, Optional
  import torch
  import torch.nn
  from torch import Tensor
label_desc: Use image node_id 2 for calling a global function with example usage: topk_sum(lprobs_list, k) and returns: topk_sum_values,  topk_sum_indices,  topk_sum_beams