def __call__(self, inputs, states):
batch_size = inputs.shape[<NUM_LIT:0>]<EOL>beam_size = self._beam_size<EOL>ctx = inputs.context<EOL>if hasattr(self._decoder, '<STR_LIT>'):<EOL><INDENT>state_info = self._decoder.state_info(batch_size)<EOL><DEDENT>else:<EOL><INDENT>state_info = None<EOL><DEDENT>states = _expand_to_beam_size(states, beam_size=beam_size, batch_size=batch_size,<EOL>state_info=state_info)<EOL>step_input = _expand_to_beam_size(inputs, beam_size=beam_size,<EOL>batch_size=batch_size).astype(np.int32)<EOL>beam_alive_mask = mx.nd.ones(shape=(batch_size, beam_size), ctx=ctx, dtype=np.int32)<EOL>valid_length = mx.nd.ones(shape=(batch_size, beam_size), ctx=ctx, dtype=np.int32)<EOL>scores = mx.nd.zeros(shape=(batch_size, beam_size), ctx=ctx)<EOL>if beam_size > <NUM_LIT:1>:<EOL><INDENT>scores[:, <NUM_LIT:1>:beam_size] = LARGE_NEGATIVE_FLOAT<EOL><DEDENT>samples = step_input.reshape((batch_size, beam_size, <NUM_LIT:1>))<EOL>for i in range(self._max_length):<EOL><INDENT>log_probs, new_states = self._decoder(step_input, states)<EOL>vocab_size_nd = mx.nd.array([log_probs.shape[<NUM_LIT:1>]], ctx=ctx, dtype=np.int32)<EOL>batch_shift_nd = mx.nd.arange(<NUM_LIT:0>, batch_size * beam_size, beam_size, ctx=ctx,<EOL>dtype=np.int32)<EOL>step_nd = mx.nd.array([i + <NUM_LIT:1>], ctx=ctx)<EOL>samples, valid_length, scores, chosen_word_ids, beam_alive_mask, states =self._updater(samples, valid_length, log_probs, scores, step_nd, beam_alive_mask,<EOL>new_states, vocab_size_nd, batch_shift_nd)<EOL>step_input = mx.nd.relu(chosen_word_ids).reshape((-<NUM_LIT:1>,))<EOL>if mx.nd.sum(beam_alive_mask).asscalar() == <NUM_LIT:0>:<EOL><INDENT>return samples, scores, valid_length<EOL><DEDENT><DEDENT>final_word = mx.nd.where(beam_alive_mask,<EOL>mx.nd.full(shape=(batch_size, beam_size),<EOL>val=self._eos_id, ctx=ctx, dtype=np.int32),<EOL>mx.nd.full(shape=(batch_size, beam_size),<EOL>val=-<NUM_LIT:1>, ctx=ctx, dtype=np.int32))<EOL>samples = mx.nd.concat(samples, final_word.reshape((<NUM_LIT:0>, <NUM_LIT:0>, <NUM_LIT:1>)), dim=<NUM_LIT:2>)<EOL>valid_length += beam_alive_mask<EOL>return samples, scores, valid_length<EOL>
Sample by beam search.

        Parameters
        ----------
        inputs : NDArray
            The initial input of the decoder. Shape is (batch_size,).
        states : Object that contains NDArrays
            The initial states of the decoder.
        Returns
        -------
        samples : NDArray
            Samples draw by beam search. Shape (batch_size, beam_size, length). dtype is int32.
        scores : NDArray
            Scores of the samples. Shape (batch_size, beam_size). We make sure that scores[i, :] are
            in descending order.
        valid_length : NDArray
            The valid length of the samples. Shape (batch_size, beam_size). dtype will be int32.