code: def test_beam_search_stochastic(early_return, eos_id):
      class SimpleStepDecoder(HybridBlock):
          def __init__(self, vocab_size=5, hidden_units=4):
              super().__init__()
              self.x2h_map = nn.Embedding(
                  input_dim=vocab_size, output_dim=hidden_units
              )
              self.h2h_map = nn.Dense(units=hidden_units, flatten=False)
              self.vocab_map = nn.Dense(units=vocab_size, flatten=False)
  
          @property
          def state_batch_axis(self):
              return 0
  
          @property
          def data_batch_axis(self):
              return 0
  
          def forward(self, data, state):
              """
  
              Parameters
              ----------
              F
              data :
                  (batch_size,)
              states :
                  (batch_size, C)
  
              Returns
              -------
              out :
                  (batch_size, vocab_size)
              new_state :
                  (batch_size, C)
              """
              new_state = self.h2h_map(state)
              out = self.vocab_map(self.x2h_map(data) + new_state)
              return out, new_state
  
      vocab_size = 3
      batch_size = 2
      hidden_units = 3
      beam_size = 4
      step_decoder = SimpleStepDecoder(vocab_size, hidden_units)
      step_decoder.initialize()
      sampler = BeamSearchSampler(
          beam_size=4,
          decoder=step_decoder,
          eos_id=eos_id,
          vocab_size=vocab_size,
          stochastic=True,
          max_length_b=100,
          early_return=early_return,
      )
      states = mx.np.random.normal(0, 1, (batch_size, hidden_units))
      inputs = mx.np.random.randint(0, vocab_size, (batch_size,))
      samples, scores, valid_length = sampler(inputs, states)
      samples = samples.asnumpy()
      valid_length = valid_length.asnumpy()
      for i in range(batch_size):
          for j in range(beam_size):
              vl = valid_length[i, j]
              if eos_id is not None:
                  assert samples[i, j, vl - 1] == eos_id
              if vl < samples.shape[2]:
                  assert (samples[i, j, vl:] == -1).all()
              assert (samples[i, :, 0] == inputs[i].asnumpy()).all()
  
      # test for repeativeness
      has_different_sample = False
      for _ in range(10):
          new_samples, scores, valid_length = sampler(inputs, states)
          if not np.array_equal(new_samples.asnumpy(), samples):
              has_different_sample = True
              break
      assert has_different_sample
  
file_path: gluon-nlp/tests/test_sequence_sampler.py
imports: import collections
  import functools
  import mxnet
  import numpy
  import scipy
  import pytest
  from mxnet.gluon import nn, HybridBlock
  from numpy.testing import assert_allclose
  from gluonnlp.sequence_sampler import BeamSearchScorer, BeamSearchSampler
label_desc: Use image node_id 3 for calling a global function with example usage: test_beam_search_stochastic(early_return, eos_id) and returns: int, int, out,  new_state