class_docstr: Batch beam search implementation.
docstr: Batch-compute topk full token ids and partial token ids.
  
  Args:
      weighted_scores (torch.Tensor): The weighted sum scores for each tokens.
          Its shape is `(n_beam, self.vocab_size)`.
      ids (torch.Tensor): The partial token ids to compute topk.
          Its shape is `(n_beam, self.pre_beam_size)`.
  
  Returns:
      Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
          The topk full (prev_hyp, new_token) ids
          and partial (prev_hyp, new_token) ids.
          Their shapes are all `(self.beam_size,)`
docstr_tok: ["Batch-compute","topk","full","token","ids","and","partial","token","ids",".","Args",":","weighted_scores","(","torch.Tensor",")",":","The","weighted","sum","scores","for","each","tokens",".","Its","shape","is","`","(","n_beam",",","self.vocab_size",")","`",".","ids","(","torch.Tensor",")",":","The","partial","token","ids","to","compute","topk",".","Its","shape","is","`","(","n_beam",",","self.pre_beam_size",")","`",".","Returns",":","Tuple","[","torch.Tensor",",","torch.Tensor",",","torch.Tensor",",","torch.Tensor","]",":","The","topk","full","(","prev_hyp",",","new_token",")","ids","and","partial","(","prev_hyp",",","new_token",")","ids",".","Their","shapes","are","all","`","(","self.beam_size",",",")","`"]
code: def batch_beam(
      self, weighted_scores: torch.Tensor, ids: torch.Tensor
  ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
      """Batch-compute topk full token ids and partial token ids.
  
      Args:
          weighted_scores (torch.Tensor): The weighted sum scores for each tokens.
              Its shape is `(n_beam, self.vocab_size)`.
          ids (torch.Tensor): The partial token ids to compute topk.
              Its shape is `(n_beam, self.pre_beam_size)`.
  
      Returns:
          Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
              The topk full (prev_hyp, new_token) ids
              and partial (prev_hyp, new_token) ids.
              Their shapes are all `(self.beam_size,)`
  
      """
      top_ids = weighted_scores.view(-1).topk(self.beam_size)[1]
      # Because of the flatten above, `top_ids` is organized as:
      # [hyp1 * V + token1, hyp2 * V + token2, ..., hypK * V + tokenK],
      # where V is `self.n_vocab` and K is `self.beam_size`
      if is_torch_1_9_plus:
          prev_hyp_ids = torch.div(
              top_ids, self.n_vocab, rounding_mode="trunc"
          )
      else:
          prev_hyp_ids = top_ids // self.n_vocab
      new_token_ids = top_ids % self.n_vocab
      return prev_hyp_ids, new_token_ids, prev_hyp_ids, new_token_ids
  
imports: import logging
  from typing import Any, Dict, List, NamedTuple, Tuple
  import torch
  from packaging.version import parse
  from torch.nn.utils.rnn import pad_sequence
  from espnet.nets.beam_search import BeamSearch, Hypothesis
label_desc: Use image node_id 5 for calling the BatchBeamSearch obj's underlying member method code with example usage: obj.batch_beam(weighted_scores, ids) and returns: prev_hyp_ids,  new_token_ids,  prev_hyp_ids,  new_token_ids