def inner_loop(i, decoded_ids, logprobs):
    """Decoder function invoked on each while loop iteration."""
    logits = symbols_to_logits_fn(decoded_ids, i)
    next_id = tf.argmax(logits, axis=0)
    softmax = tf.nn.softmax(logits)
    extended_vocab_size = tf.shape(softmax)[-1]
    mask = tf.one_hot(next_id, extended_vocab_size)
    prob = tf.reduce_sum(softmax * mask)
    logprob = tf.log(prob)

    # Add one-hot values to output Tensors, since values at index > i+1 should
    # still be zero.
    logprobs += tf.one_hot(
        i + 1, decode_length + 1, on_value=logprob, dtype=tf.float32)
    decoded_ids += tf.one_hot(
        i + 1, decode_length + 1, on_value=next_id, dtype=tf.int64)

    return i + 1, decoded_ids, logprobs
Decoder function invoked on each while loop iteration.
language/xsp/model/transformer.py