docstr: Parameters
  ----------
  samples : mx.np.ndarray
      The current samples generated by beam search.
      Shape (batch_size, beam_size, L).
  valid_length : mx.np.ndarray
      The current valid lengths of the samples
  outputs : mx.np.ndarray
      Outputs from predictor. If from_logits was set to True in scorer, then it's the
      log probability of the current step. Else, it's the unnormalized outputs before
      softmax or log_softmax.
      Shape (batch_size * beam_size, V).
  scores : mx.np.ndarray
      The previous scores. Shape (batch_size, beam_size)
  step : mx.np.ndarray
      The current step for doing beam search. Begins from 1. Shape ()
  beam_alive_mask : mx.np.ndarray
      Shape (batch_size, beam_size)
  states : nested structure of mx.np.ndarray
      Each mx.np.ndarray should have shape (N, ...) when state_info is None,
      or same as the layout in state_info when it's not None.
  batch_shift : mx.np.ndarray
      Contains [0, beam_size, 2 * beam_size, ..., (batch_size - 1) * beam_size].
      Shape (batch_size,)
  
  Returns
  -------
  new_samples : mx.np.ndarray or an empty list
      The updated samples.
      When single_step is False, shape (batch_size, beam_size, L + 1)
  new_valid_length : mx.np.ndarray
      Valid lengths of the samples. Shape (batch_size, beam_size)
  new_scores : mx.np.ndarray
      Shape (batch_size, beam_size)
  chosen_word_ids : mx.np.ndarray
      The chosen word ids of the step. Shape (batch_size, beam_size). If it's negative,
      no word will be appended to the beam.
  beam_alive_mask : mx.np.ndarray
      Shape (batch_size, beam_size)
  new_states : nested structure of mx.np.ndarray
      Inner mx.np.ndarrays have shape (batch_size * beam_size, ...)
docstr_tok: ["Parameters","--","--","--","--","--","samples",":","mx.np.ndarray","The","current","samples","generated","by","beam","search",".","Shape","(","batch_size",",","beam_size",",","L",")",".","valid_length",":","mx.np.ndarray","The","current","valid","lengths","of","the","samples","outputs",":","mx.np.ndarray","Outputs","from","predictor",".","If","from_logits","was","set","to","True","in","scorer",",","then","it","'s","the","log","probability","of","the","current","step",".","Else",",","it","'s","the","unnormalized","outputs","before","softmax","or","log_softmax",".","Shape","(","batch_size","*","beam_size",",","V",")",".","scores",":","mx.np.ndarray","The","previous","scores",".","Shape","(","batch_size",",","beam_size",")","step",":","mx.np.ndarray","The","current","step","for","doing","beam","search",".","Begins","from","1",".","Shape","(",")","beam_alive_mask",":","mx.np.ndarray","Shape","(","batch_size",",","beam_size",")","states",":","nested","structure","of","mx.np.ndarray","Each","mx.np.ndarray","should","have","shape","(","N",",","...",")","when","state_info","is","None",",","or","same","as","the","layout","in","state_info","when","it","'s","not","None",".","batch_shift",":","mx.np.ndarray","Contains","[","0",",","beam_size",",","2","*","beam_size",",","...",",","(","batch_size","-","1",")","*","beam_size","]",".","Shape","(","batch_size",",",")","Returns","--","--","--","-","new_samples",":","mx.np.ndarray","or","an","empty","list","The","updated","samples",".","When","single_step","is","False",",","shape","(","batch_size",",","beam_size",",","L","+","1",")","new_valid_length",":","mx.np.ndarray","Valid","lengths","of","the","samples",".","Shape","(","batch_size",",","beam_size",")","new_scores",":","mx.np.ndarray","Shape","(","batch_size",",","beam_size",")","chosen_word_ids",":","mx.np.ndarray","The","chosen","word","ids","of","the","step",".","Shape","(","batch_size",",","beam_size",")",".","If","it","'s","negative",",","no","word","will","be","appended","to","the","beam",".","beam_alive_mask",":","mx.np.ndarray","Shape","(","batch_size",",","beam_size",")","new_states",":","nested","structure","of","mx.np.ndarray","Inner","mx.np.ndarrays","have","shape","(","batch_size","*","beam_size",",","...",")"]
code: def forward(
      self,
      samples,
      valid_length,
      outputs,
      scores,
      step,
      beam_alive_mask,  # pylint: disable=arguments-differ
      states,
      batch_shift,
  ):
      """
  
      Parameters
      ----------
      samples : mx.np.ndarray
          The current samples generated by beam search.
          Shape (batch_size, beam_size, L).
      valid_length : mx.np.ndarray
          The current valid lengths of the samples
      outputs : mx.np.ndarray
          Outputs from predictor. If from_logits was set to True in scorer, then it's the
          log probability of the current step. Else, it's the unnormalized outputs before
          softmax or log_softmax.
          Shape (batch_size * beam_size, V).
      scores : mx.np.ndarray
          The previous scores. Shape (batch_size, beam_size)
      step : mx.np.ndarray
          The current step for doing beam search. Begins from 1. Shape ()
      beam_alive_mask : mx.np.ndarray
          Shape (batch_size, beam_size)
      states : nested structure of mx.np.ndarray
          Each mx.np.ndarray should have shape (N, ...) when state_info is None,
          or same as the layout in state_info when it's not None.
      batch_shift : mx.np.ndarray
          Contains [0, beam_size, 2 * beam_size, ..., (batch_size - 1) * beam_size].
          Shape (batch_size,)
  
      Returns
      -------
      new_samples : mx.np.ndarray or an empty list
          The updated samples.
          When single_step is False, shape (batch_size, beam_size, L + 1)
      new_valid_length : mx.np.ndarray
          Valid lengths of the samples. Shape (batch_size, beam_size)
      new_scores : mx.np.ndarray
          Shape (batch_size, beam_size)
      chosen_word_ids : mx.np.ndarray
          The chosen word ids of the step. Shape (batch_size, beam_size). If it's negative,
          no word will be appended to the beam.
      beam_alive_mask : mx.np.ndarray
          Shape (batch_size, beam_size)
      new_states : nested structure of mx.np.ndarray
          Inner mx.np.ndarrays have shape (batch_size * beam_size, ...)
      """
      beam_size = self._beam_size
      vocab_size = self._vocab_size
      beam_alive_mask_bcast = mx.np.expand_dims(beam_alive_mask, axis=2)
      candidate_scores = self._scorer(
          mx.npx.reshape(outputs, (-6, -1, beam_size, -2)), scores, step
      )
      if self.stochastic:
          if step == 1:
              candidate_scores_gumbel = candidate_scores[
                  :1
              ] + mx.np.random.gumbel(
                  mx.np.zeros_like(candidate_scores[:1])
              )
              candidate_scores_residual = candidate_scores[1:]
              candidate_scores = mx.np.concatenate(
                  (candidate_scores_gumbel, candidate_scores_residual),
                  axis=0,
              )
          else:
              candidate_scores = self.gumbel_with_maximum(
                  candidate_scores, scores, -1
              )
      # Concat the candidate scores and the scores of the finished beams
      # The resulting candidate score will have shape (batch_size, beam_size * |V| + beam_size)
      candidate_scores = mx.np.where(
          beam_alive_mask_bcast,
          candidate_scores,
          mx.np.full_like(candidate_scores, LARGE_NEGATIVE_FLOAT),
      )
      finished_scores = mx.np.where(
          beam_alive_mask,
          mx.np.full_like(scores, LARGE_NEGATIVE_FLOAT),
          scores,
      )
      candidate_scores = mx.np.concatenate(
          [mx.npx.reshape(candidate_scores, (-2, -1)), finished_scores],
          axis=1,
      )
      # Get the top K scores
      # new_scores and indices will have shape (batch_size, beam_size)
      new_scores, indices = mx.npx.topk(
          candidate_scores, axis=1, k=beam_size, ret_typ="both"
      )
      indices = indices.astype(mx.np.int32)
      use_prev = (indices >= (beam_size * vocab_size)).astype(
          mx.np.int32
      )
      chosen_word_ids = mx.np.mod(indices, vocab_size)
      beam_ids = mx.np.where(
          use_prev,
          indices - beam_size * vocab_size,
          mx.np.floor(indices / vocab_size).astype(mx.np.int32),
      )
      batch_beam_indices = beam_ids + mx.np.expand_dims(
          batch_shift, axis=1
      )
      chosen_word_ids = mx.np.where(
          use_prev, -mx.np.ones_like(indices), chosen_word_ids
      )
      # Update the samples and vaild_length
      # TODO(sxjscience) The current implementation is quite tricky
      #  We should wait for hybridizable advanced indexing to avoid this
      selected_samples = mx.np.take(
          mx.npx.reshape(samples, (-5, -2)),
          batch_beam_indices.reshape((-1,)),
          axis=0,
      )
      new_samples = mx.npx.reshape(
          mx.np.concatenate(
              [selected_samples, chosen_word_ids.reshape((-1, 1))],
              axis=1,
          ),
          (-6, -1, beam_size, -2),
      )
      new_valid_length = (
          mx.np.take(
              valid_length.reshape((-1,)),
              batch_beam_indices.reshape((-1,)),
              axis=0,
          ).reshape((-1, beam_size))
          + 1
          - use_prev
      )
      # Update the states
      new_states = _choose_states(
          states,
          batch_beam_indices.reshape((-1,)),
          self._state_batch_axis,
      )
      # Update the alive mask.
      beam_alive_mask = mx.np.take(
          beam_alive_mask.reshape((-1,)),
          batch_beam_indices.reshape((-1,)),
          axis=0,
      ).reshape((-1, beam_size))
      if self._eos_id is not None:
          beam_alive_mask = beam_alive_mask * (
              chosen_word_ids != self._eos_id
          ).astype(mx.np.float32)
      return (
          new_samples,
          new_valid_length,
          new_scores,
          chosen_word_ids,
          beam_alive_mask,
          new_states,
      )
  
imports: import warnings
  import mxnet
  import abc
  from mxnet.gluon import HybridBlock
  from typing import Callable, Optional
  from .layers import get_activation
label_desc: Use image node_id 4 for calling the _BeamSearchStepUpdate obj's underlying member method code with example usage: obj.forward(samples, valid_length, outputs, scores, step, beam_alive_mask, states, batch_shift) and returns: new_samples,  new_valid_length,  new_scores,  chosen_word_ids,  beam_alive_mask,  new_states