docstr: Take a single search step.
  
  Args:
      step: the current search step, starting at 0
      lprobs: dictionary of channels {channel : (bsz x input_beam_size x vocab_size_channel)}
          the model's log-probabilities over the vocabulary at the current step
      scores: {channel : (bsz x input_beam_size x step)}
          the historical model scores of each hypothesis up to this point
      prev_output_tokens: {channel : (bsz x step)}
          the previously generated oputput tokens
      original_batch_idxs: (bsz)
          the tensor with the batch indices, in the range [0, bsz)
          this is useful in case there has been applied a re-ordering
          and we need to know the orignal indices
  
  Return: A tuple of (scores, indices, beams) where:
      scores: {channel : (bsz x output_beam_size)}
          the scores of the chosen elements; output_beam_size can be
          larger than input_beam_size, e.g., we may return
          2*input_beam_size to account for EOS
      indices: {channel : (bsz x output_beam_size)}
          the indices of the chosen elements
      beams: (bsz x output_beam_size)
          the hypothesis ids of the chosen elements, in the range [0, input_beam_size)
docstr_tok: ["Take","a","single","search","step",".","Args",":","step",":","the","current","search","step",",","starting","at","0","lprobs",":","dictionary","of","channels","{","channel",":","(","bsz","x","input_beam_size","x","vocab_size_channel",")","}","the","model","'s","log-probabilities","over","the","vocabulary","at","the","current","step","scores",":","{","channel",":","(","bsz","x","input_beam_size","x","step",")","}","the","historical","model","scores","of","each","hypothesis","up","to","this","point","prev_output_tokens",":","{","channel",":","(","bsz","x","step",")","}","the","previously","generated","oputput","tokens","original_batch_idxs",":","(","bsz",")","the","tensor","with","the","batch","indices",",","in","the","range","[","0",",","bsz",")","this","is","useful","in","case","there","has","been","applied","a","re-ordering","and","we","need","to","know","the","orignal","indices","Return",":","A","tuple","of","(","scores",",","indices",",","beams",")","where",":","scores",":","{","channel",":","(","bsz","x","output_beam_size",")","}","the","scores","of","the","chosen","elements",";","output_beam_size","can","be","larger","than","input_beam_size",",","e.g.",",","we","may","return","2","*","input_beam_size","to","account","for","EOS","indices",":","{","channel",":","(","bsz","x","output_beam_size",")","}","the","indices","of","the","chosen","elements","beams",":","(","bsz","x","output_beam_size",")","the","hypothesis","ids","of","the","chosen","elements",",","in","the","range","[","0",",","input_beam_size",")"]
code: def step(
      self,
      step,
      lprobs,
      scores,
      prev_output_tokens=None,
      original_batch_idxs=None,
  ):
      """Take a single search step.
  
      Args:
          step: the current search step, starting at 0
          lprobs: dictionary of channels {channel : (bsz x input_beam_size x vocab_size_channel)}
              the model's log-probabilities over the vocabulary at the current step
          scores: {channel : (bsz x input_beam_size x step)}
              the historical model scores of each hypothesis up to this point
          prev_output_tokens: {channel : (bsz x step)}
              the previously generated oputput tokens
          original_batch_idxs: (bsz)
              the tensor with the batch indices, in the range [0, bsz)
              this is useful in case there has been applied a re-ordering
              and we need to know the orignal indices
  
      Return: A tuple of (scores, indices, beams) where:
          scores: {channel : (bsz x output_beam_size)}
              the scores of the chosen elements; output_beam_size can be
              larger than input_beam_size, e.g., we may return
              2*input_beam_size to account for EOS
          indices: {channel : (bsz x output_beam_size)}
              the indices of the chosen elements
          beams: (bsz x output_beam_size)
              the hypothesis ids of the chosen elements, in the range [0, input_beam_size)
      """
      raise NotImplementedError
  
imports: from typing import Dict, Optional
  import torch
  import torch.nn
  from torch import Tensor
label_desc: Use image node_id 2 for calling the MultichannelSearch obj's underlying member method code with example usage: obj.step(step, lprobs, scores, prev_output_tokens, original_batch_idxs) without return types