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,
      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, ...)
      """
      # bsz * beam_size * vocab_size
      outputs = outputs.reshape((-1, self._beam_size, self._vocab_size))
      probs = mx.npx.softmax(outputs / self._temperature)
  
      if self._sampling_topp > 0:
          probs = mx.np.where(
              probs > self._sampling_topp,
              probs,
              mx.np.zeros_like(probs),
          )
      elif self._sampling_topk > 0:
          topk_probs = mx.npx.topk(
              probs, axis=2, k=self._sampling_topk, ret_typ="value"
          )
          # choose the k max prob
          k_prob = topk_probs[:, :, -1]
          k_prob = mx.np.expand_dims(k_prob, axis=-1)
          probs = mx.np.where(
              probs >= k_prob, probs, mx.np.zeros_like(probs)
          )
  
      # renormalize
      probs_sum = mx.np.sum(probs, axis=2, keepdims=True)
      probs = probs / probs_sum
  
      # bsz * beam_size
      (
          chosen_word_ids,
          chosen_word_log_probs,
      ) = mx.npx.random.categorical(probs, get_prob=True)
  
      new_scores = scores + mx.np.where(
          beam_alive_mask,
          chosen_word_log_probs,
          mx.np.zeros_like(chosen_word_log_probs),
      )
  
      # mask dead words
      chosen_word_ids = mx.np.where(
          beam_alive_mask,
          chosen_word_ids,
          mx.np.full_like(beam_alive_mask, -1, dtype=mx.np.int32),
      )
  
      new_valid_length = valid_length + beam_alive_mask.astype(
          mx.np.int32
      )
      new_samples = mx.np.concatenate(
          [samples, mx.np.expand_dims(chosen_word_ids, axis=2)], axis=2
      )
      new_states = states
      if self._eos_id is not None:
          beam_alive_mask = beam_alive_mask * (
              chosen_word_ids != self._eos_id
          ).astype(mx.np.int32)
  
      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 2 for calling the _MultinomialStepUpdate 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