class_docstr: Batch beam search implementation.
docstr: Search new tokens for running hypotheses and encoded speech x.
  
  Args:
      running_hyps (BatchHypothesis): Running hypotheses on beam
      x (torch.Tensor): Encoded speech feature (T, D)
      pre_x (torch.Tensor): Encoded speech feature for sequential attention (T, D)
  
  Returns:
      BatchHypothesis: Best sorted hypotheses
docstr_tok: ["Search","new","tokens","for","running","hypotheses","and","encoded","speech","x.","Args",":","running_hyps","(","BatchHypothesis",")",":","Running","hypotheses","on","beam","x","(","torch.Tensor",")",":","Encoded","speech","feature","(","T",",","D",")","pre_x","(","torch.Tensor",")",":","Encoded","speech","feature","for","sequential","attention","(","T",",","D",")","Returns",":","BatchHypothesis",":","Best","sorted","hypotheses"]
code: def search(
      self,
      running_hyps: BatchHypothesis,
      x: torch.Tensor,
      pre_x: torch.Tensor = None,
  ) -> BatchHypothesis:
      """Search new tokens for running hypotheses and encoded speech x.
  
      Args:
          running_hyps (BatchHypothesis): Running hypotheses on beam
          x (torch.Tensor): Encoded speech feature (T, D)
          pre_x (torch.Tensor): Encoded speech feature for sequential attention (T, D)
  
      Returns:
          BatchHypothesis: Best sorted hypotheses
  
      """
      n_batch = len(running_hyps)
      part_ids = None  # no pre-beam
      # batch scoring
      weighted_scores = torch.zeros(
          n_batch, self.n_vocab, dtype=x.dtype, device=x.device
      )
      if self.return_hs:
          hs, scores, states = self.score_full(
              running_hyps,
              x.expand(n_batch, *x.shape),
              pre_x=pre_x.expand(n_batch, *pre_x.shape)
              if pre_x is not None
              else None,
          )
      else:
          scores, states = self.score_full(
              running_hyps,
              x.expand(n_batch, *x.shape),
              pre_x=pre_x.expand(n_batch, *pre_x.shape)
              if pre_x is not None
              else None,
          )
  
      for k in self.full_scorers:
          weighted_scores += self.weights[k] * scores[k]
      # partial scoring
      if self.do_pre_beam:
          pre_beam_scores = (
              weighted_scores
              if self.pre_beam_score_key == "full"
              else scores[self.pre_beam_score_key]
          )
          part_ids = torch.topk(
              pre_beam_scores, self.pre_beam_size, dim=-1
          )[1]
      # NOTE(takaaki-hori): Unlike BeamSearch, we assume that score_partial returns
      # full-size score matrices, which has non-zero scores for part_ids and zeros
      # for others.
      part_scores, part_states = self.score_partial(
          running_hyps, part_ids, x, pre_x
      )
      for k in self.part_scorers:
          weighted_scores += self.weights[k] * part_scores[k]
      # add previous hyp scores
      weighted_scores += running_hyps.score.to(
          dtype=x.dtype, device=x.device
      ).unsqueeze(1)
  
      # TODO(karita): do not use list. use batch instead
      # see also https://github.com/espnet/espnet/pull/1402#discussion_r354561029
      # update hyps
      best_hyps = []
      prev_hyps = self.unbatchfy(running_hyps)
      for (
          full_prev_hyp_id,
          full_new_token_id,
          part_prev_hyp_id,
          part_new_token_id,
      ) in zip(*self.batch_beam(weighted_scores, part_ids)):
          prev_hyp = prev_hyps[full_prev_hyp_id]
          if self.return_hs:
              new_hs = prev_hyp.hs + [hs[full_prev_hyp_id].squeeze(0)]
          else:
              new_hs = []
          best_hyps.append(
              Hypothesis(
                  score=weighted_scores[
                      full_prev_hyp_id, full_new_token_id
                  ],
                  yseq=self.append_token(
                      prev_hyp.yseq, full_new_token_id
                  ),
                  scores=self.merge_scores(
                      prev_hyp.scores,
                      {
                          k: v[full_prev_hyp_id]
                          for k, v in scores.items()
                      },
                      full_new_token_id,
                      {
                          k: v[part_prev_hyp_id]
                          for k, v in part_scores.items()
                      },
                      part_new_token_id,
                  ),
                  states=self.merge_states(
                      {
                          k: self.full_scorers[k].select_state(
                              v, full_prev_hyp_id
                          )
                          for k, v in states.items()
                      },
                      {
                          k: self.part_scorers[k].select_state(
                              v, part_prev_hyp_id, part_new_token_id
                          )
                          for k, v in part_states.items()
                      },
                      part_new_token_id,
                  ),
                  hs=new_hs,
              )
          )
      return self.batchfy(best_hyps)
  
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 10 for calling the BatchBeamSearch obj's underlying member method code with example usage: obj.search(running_hyps, x, pre_x) and returns: self