'Beam search decoding.
Args:
features: an map of string to `Tensor`
decode_length: an integer.  How many additional timesteps to decode.
beam_size: number of beams.
top_beams: an integer. How many of the beams to return.
last_position_only: a boolean, speed-up by computing last position only.
alpha: Float that controls the length penalty. larger the alpha, stronger
the preference for slonger translations.
Returns:
samples: an integer `Tensor`. Top samples from the beam search'

def _beam_decode(self, features, decode_length, beam_size, top_beams, last_position_only, alpha):

    def symbols_to_logits_fn(ids):
        'Go    from    ids    to    logits.'
        ids = tf.expand_dims(tf.expand_dims(ids, axis=2), axis=3)
        ids = tf.pad(ids[:, 1:], [[0, 0], [0, 1], [0, 0], [0, 0]])
        features['targets'] = ids
        self._coverage = None
        (sharded_logits, _) = self.model_fn(features, False, last_position_only=last_position_only)
        logits = sharded_logits[0]
        if last_position_only:
            return tf.squeeze(logits, axis=[1, 2, 3])
        current_output_position = (tf.shape(ids)[1] - 1)
        logits = logits[:, current_output_position, :, :]
        return tf.squeeze(logits, axis=[1, 2])
    batch_size = tf.shape(features['inputs'])[0]
    initial_ids = tf.zeros([batch_size], dtype=tf.int32)
    inputs_old = features['inputs']
    features['inputs'] = tf.expand_dims(features['inputs'], 1)
    if (len(features['inputs'].shape) < 5):
        features['inputs'] = tf.expand_dims(features['inputs'], 4)
    features['inputs'] = tf.tile(features['inputs'], [1, beam_size, 1, 1, 1])
    s = tf.shape(features['inputs'])
    features['inputs'] = tf.reshape(features['inputs'], [(s[0] * s[1]), s[2], s[3], s[4]])
    target_modality = self._hparams.problems[self._problem_idx].target_modality
    vocab_size = target_modality.top_dimensionality
    decode_length = (tf.shape(features['inputs'])[1] + tf.constant(decode_length))
    (ids, scores) = beam_search.beam_search(symbols_to_logits_fn, initial_ids, beam_size, decode_length, vocab_size, alpha)
    features['inputs'] = inputs_old
    return_scores = False
    if (top_beams == 1):
        if return_scores:
            return {'outputs': ids[:, 0, 1:], 'scores': scores}
        return ids[:, 0, 1:]
    else:
        if return_scores:
            return {'outputs': ids[:, :top_beams, 1:], 'scores': scores}
        return ids[:, :top_beams, 1:]
