class_docstr: Draw samples from the decoder by beam search.
  
  Parameters
  ----------
  beam_size : int
      The beam size.
  decoder
      Function of the one-step-ahead decoder, should inherit SequenceDecoder and
      have the form::
  
          outputs, new_states = decoder(step_input, states)
  
      The outputs, input should follow these rules:
  
      - step_input has shape (batch_size,),
      - outputs has shape (batch_size, V),
      - states and new_states have the same structure.
  eos_id
      Id of the EOS token. No other elements will be appended to the sample if it reaches eos_id,
      if eos_id == None, beam will not die
  scorer : BeamSearchScorer, default BeamSearchScorer(alpha=1.0, K=5)
      The score function used in beam search.
  max_length_a
      The `a` value in the formula `a * x + b`. Generate sequences of maximum length `a * x + b`,
      where `x` is the maximum source length.
  max_length_b
      The b value in the formula `a * x + b`. Generate sequences of maximum length `a * x + b`,
      where `x` is the maximum source length.
  min_length
      The minimum length of the generated sequences.
  temperature
      The temperature.
  stochastic
      Whether to use stochastic beam search,
      see [ICML2019] "Stochastic Beams and Where to Find Them" for detail.
      https://arxiv.org/abs/1903.06059
  sampling
      Whether to use multinomial sampler.
  sampling_topp
      Multinomial sampling with topp,
      see [ICLR2020] "The Curious Case of Neural Text Degeneration"
      https://arxiv.org/abs/1904.09751
  sampling_topk
      Multinomial sampling with topk,
      see [ACL2018] "Hierarchical Neural Story Generation"
      https://www.aclweb.org/anthology/P18-1082.pdf
  early_return
      Whether to return when all beams are dead.
      Without early_return, the sequences will be generated until the
      maximum length is reached.
class_docstr_tok: ["Draw","samples","from","the","decoder","by","beam","search",".","Parameters","--","--","--","--","--","beam_size",":","int","The","beam","size",".","decoder","Function","of","the","one-step-ahead","decoder",",","should","inherit","SequenceDecoder","and","have","the","form",":",":","outputs",",","new_states","=","decoder","(","step_input",",","states",")","The","outputs",",","input","should","follow","these","rules",":","-","step_input","has","shape","(","batch_size",",",")",",","-","outputs","has","shape","(","batch_size",",","V",")",",","-","states","and","new_states","have","the","same","structure",".","eos_id","Id","of","the","EOS","token",".","No","other","elements","will","be","appended","to","the","sample","if","it","reaches","eos_id",",","if","eos_id","==","None",",","beam","will","not","die","scorer",":","BeamSearchScorer",",","default","BeamSearchScorer","(","alpha=1.0",",","K=5",")","The","score","function","used","in","beam","search",".","max_length_a","The","`","a","`","value","in","the","formula","`","a","*","x","+","b","`",".","Generate","sequences","of","maximum","length","`","a","*","x","+","b","`",",","where","`","x","`","is","the","maximum","source","length",".","max_length_b","The","b","value","in","the","formula","`","a","*","x","+","b","`",".","Generate","sequences","of","maximum","length","`","a","*","x","+","b","`",",","where","`","x","`","is","the","maximum","source","length",".","min_length","The","minimum","length","of","the","generated","sequences",".","temperature","The","temperature",".","stochastic","Whether","to","use","stochastic","beam","search",",","see","[","ICML2019","]","``","Stochastic","Beams","and","Where","to","Find","Them","''","for","detail",".","https",":","\/\/arxiv.org\/abs\/1903.06059","sampling","Whether","to","use","multinomial","sampler",".","sampling_topp","Multinomial","sampling","with","topp",",","see","[","ICLR2020","]","``","The","Curious","Case","of","Neural","Text","Degeneration","''","https",":","\/\/arxiv.org\/abs\/1904.09751","sampling_topk","Multinomial","sampling","with","topk",",","see","[","ACL2018","]","``","Hierarchical","Neural","Story","Generation","''","https",":","\/\/www.aclweb.org\/anthology\/P18-1082.pdf","early_return","Whether","to","return","when","all","beams","are","dead",".","Without","early_return",",","the","sequences","will","be","generated","until","the","maximum","length","is","reached","."]
docstr: Sample by beam search.
  
  Parameters
  ----------
  inputs : mx.np.ndarray
      The initial input of the decoder. Shape is (batch_size,).
  states : Object that contains mx.np.ndarrays
      The initial states of the decoder.
  src_seq_lengths : mx.np.ndarray
      The source sequence lengths. Shape is (batch_size,).
  
  Returns
  -------
  samples : mx.np.ndarray
      Samples draw by beam search. Shape (batch_size, beam_size, length).
      DType is int32.
  scores : mx.np.ndarray
      Scores of the samples. Shape (batch_size, beam_size).
       We make sure that scores[i, :] are in descending order.
  valid_length : mx.np.ndarray
      The valid length of the samples. Shape (batch_size, beam_size).
      DType is int32.
docstr_tok: ["Sample","by","beam","search",".","Parameters","--","--","--","--","--","inputs",":","mx.np.ndarray","The","initial","input","of","the","decoder",".","Shape","is","(","batch_size",",",")",".","states",":","Object","that","contains","mx.np.ndarrays","The","initial","states","of","the","decoder",".","src_seq_lengths",":","mx.np.ndarray","The","source","sequence","lengths",".","Shape","is","(","batch_size",",",")",".","Returns","--","--","--","-","samples",":","mx.np.ndarray","Samples","draw","by","beam","search",".","Shape","(","batch_size",",","beam_size",",","length",")",".","DType","is","int32",".","scores",":","mx.np.ndarray","Scores","of","the","samples",".","Shape","(","batch_size",",","beam_size",")",".","We","make","sure","that","scores","[","i",",",":","]","are","in","descending","order",".","valid_length",":","mx.np.ndarray","The","valid","length","of","the","samples",".","Shape","(","batch_size",",","beam_size",")",".","DType","is","int32","."]
code: def forward(self, inputs, states, src_seq_lengths=None):
      """Sample by beam search.
  
      Parameters
      ----------
      inputs : mx.np.ndarray
          The initial input of the decoder. Shape is (batch_size,).
      states : Object that contains mx.np.ndarrays
          The initial states of the decoder.
      src_seq_lengths : mx.np.ndarray
          The source sequence lengths. Shape is (batch_size,).
  
      Returns
      -------
      samples : mx.np.ndarray
          Samples draw by beam search. Shape (batch_size, beam_size, length).
          DType is int32.
      scores : mx.np.ndarray
          Scores of the samples. Shape (batch_size, beam_size).
           We make sure that scores[i, :] are in descending order.
      valid_length : mx.np.ndarray
          The valid length of the samples. Shape (batch_size, beam_size).
          DType is int32.
      """
      ctx = inputs.ctx
      batch_size = inputs.shape[self._data_batch_axis]
      beam_size = self._beam_size
      if src_seq_lengths is not None:
          max_src_sequence_length = int(src_seq_lengths.asnumpy().max())
          max_length = max(
              self._min_length,
              max_src_sequence_length * self._max_length_a
              + self._max_length_b,
          )
      else:
          if self._max_length_a != 0:
              raise ValueError(
                  "If src_seq_lengths is not given, max_length_a must be 0!"
                  " Received {}".format(self._max_length_a)
              )
          max_length = max(self._min_length, self._max_length_b)
      # Tile the states and inputs to have shape (batch_size * beam_size, ...)
      states = _expand_to_beam_size(
          states,
          beam_size=beam_size,
          batch_size=batch_size,
          state_batch_axis=self._state_batch_axis,
      )
      step_input = _expand_to_beam_size(
          inputs,
          beam_size=beam_size,
          batch_size=batch_size,
          state_batch_axis=self._data_batch_axis,
      ).astype(mx.np.int32)
      # All beams are initialized to alive
      # Generated samples are initialized to be the inputs
      # Except the first beam where the scores are set to be zero, all beams have -inf scores.
      # Valid length is initialized to be 1
      beam_alive_mask = mx.np.ones(
          shape=(batch_size, beam_size), ctx=ctx, dtype=mx.np.float32
      )
      valid_length = mx.np.ones(
          shape=(batch_size, beam_size), ctx=ctx, dtype=mx.np.int32
      )
      scores = mx.np.zeros(shape=(batch_size, beam_size), ctx=ctx)
      if beam_size > 1:
          scores[:, 1:beam_size] = LARGE_NEGATIVE_FLOAT
      samples = step_input.reshape((batch_size, beam_size, -1))
      batch_shift = mx.np.arange(
          0,
          batch_size * beam_size,
          beam_size,
          ctx=ctx,
          dtype=mx.np.int32,
      )
      step = mx.np.array(0, ctx=ctx, dtype=mx.np.float32)
      for i in range(max_length):
          log_probs, new_states = self._decoder(step_input, states)
          assert log_probs.shape[1] == self._vocab_size
          step = step + 1
          (
              samples,
              valid_length,
              scores,
              chosen_word_ids,
              beam_alive_mask,
              states,
          ) = self._updater(
              samples,
              valid_length,
              log_probs,
              scores,
              step,
              beam_alive_mask,
              new_states,
              batch_shift,
          )
          step_input = mx.npx.relu(chosen_word_ids).reshape((-1,))
          if self._early_return:
              if mx.np.sum(beam_alive_mask).asnumpy() == 0:
                  return samples, scores, valid_length
      beam_alive_mask = beam_alive_mask.astype(mx.np.int32)
      if self._eos_id is not None:
          final_word = mx.np.where(
              beam_alive_mask,
              mx.np.full(
                  (batch_size, beam_size),
                  self._eos_id,
                  ctx=ctx,
                  dtype=mx.np.int32,
              ),
              mx.np.full(
                  (batch_size, beam_size),
                  -1,
                  ctx=ctx,
                  dtype=mx.np.int32,
              ),
          )
          samples = mx.np.concatenate(
              [
                  samples,
                  final_word.reshape(
                      (final_word.shape[0], final_word.shape[1], 1)
                  ),
              ],
              axis=2,
          )
          valid_length = valid_length + beam_alive_mask
      return samples, scores, valid_length
  
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 BeamSearchSampler obj's underlying member method code with example usage: obj.forward(inputs, states, src_seq_lengths) and returns: samples,  scores,  valid_length, samples,  scores,  valid_length