def _get_symbols_to_logits_fn(self, max_decode_length, training):
    'Returns a decoding function that calculates logits of the next tokens.'
    timing_signal = self._position_embedding_layer(inputs=None, length=max_decode_length)
    decoder_self_attention_bias = model_utils.get_decoder_self_attention_bias(max_decode_length)

    def symbols_to_logits_fn(ids, i, cache):
        'Generate logits for next potential IDs.\n\n      Args:\n        ids: Current decoded sequences. int tensor with shape [batch_size *\n          beam_size, i + 1].\n        i: Loop index.\n        cache: dictionary of values storing the encoder output, encoder-decoder\n          attention bias, and previous decoder attention values.\n\n      Returns:\n        Tuple of\n          (logits with shape [batch_size * beam_size, vocab_size],\n           updated cache values)\n      '
        decoder_input = ids[:, (- 1):]
        decoder_input = self._word_embedding_layer(decoder_input)
        decoder_input += timing_signal[i:(i + 1)]
        self_attention_bias = decoder_self_attention_bias[:, :, i:(i + 1), :(i + 1)]
        decoder_outputs = self._decoder(decoder_input, cache.get('encoder_outputs'), self_attention_bias, cache.get('encoder_decoder_attention_bias'), training=training, cache=cache)
        decoder_outputs = decoder_outputs[:, (- 1), :]
        logits = self._word_layer(decoder_outputs)
        return (logits, cache)
    return symbols_to_logits_fn
Returns a decoding function that calculates logits of the next tokens.
def _get_symbols_to_logits_fn(self, max_decode_length, training):
    
    timing_signal = self._position_embedding_layer(inputs=None, length=max_decode_length)
    decoder_self_attention_bias = model_utils.get_decoder_self_attention_bias(max_decode_length)

    def symbols_to_logits_fn(ids, i, cache):
        'Generate logits for next potential IDs.\n\n      Args:\n        ids: Current decoded sequences. int tensor with shape [batch_size *\n          beam_size, i + 1].\n        i: Loop index.\n        cache: dictionary of values storing the encoder output, encoder-decoder\n          attention bias, and previous decoder attention values.\n\n      Returns:\n        Tuple of\n          (logits with shape [batch_size * beam_size, vocab_size],\n           updated cache values)\n      '
        decoder_input = ids[:, (- 1):]
        decoder_input = self._word_embedding_layer(decoder_input)
        decoder_input += timing_signal[i:(i + 1)]
        self_attention_bias = decoder_self_attention_bias[:, :, i:(i + 1), :(i + 1)]
        decoder_outputs = self._decoder(decoder_input, cache.get('encoder_outputs'), self_attention_bias, cache.get('encoder_decoder_attention_bias'), training=training, cache=cache)
        decoder_outputs = decoder_outputs[:, (- 1), :]
        logits = self._word_layer(decoder_outputs)
        return (logits, cache)
    return symbols_to_logits_fn