def symbols_to_logits_fn(input_ids, i, state):
    """Go from ids to logits for next symbol."""
    # Size of expanded tensor (expanded by beam size).
    batch_size = tensor_utils.shape(input_ids, 0)

    # [batch_size, 1]
    current_step_mask = tf.ones([batch_size, 1], tf.int32)

    # [batch_size, num_steps]
    written_mask = tf.cast(tf.less(tf.range(num_steps), i), tf.int32)
    written_mask = tf.tile(tf.expand_dims(written_mask, 0), [batch_size, 1])
    is_written = tf.cast(written_mask, tf.bool)

    # [batch_size, cache_size + num_steps, num_layers, num_heads, head_size]
    input_cache = TransformerCache(
        keys=tf.concat([state.encoder_cache.keys, state.output_cache.keys], 1),
        values=tf.concat(
            [state.encoder_cache.values, state.output_cache.values], 1))

    # [batch_size, 1, cache_size + num_steps]
    masks = [state.encoder_cache_mask, written_mask, current_step_mask]
    attention_mask = tf.concat(masks, axis=1)
    attention_mask = tf.expand_dims(attention_mask, 1)

    # sequence_output: [batch_size, 1, hidden_size],
    # step_cache: [batch_size, 1, num_layers, num_heads, head_size]
    sequence_output, step_cache = model.compute_transformer(
        input_ids=input_ids,
        input_segment_id=tf.fill(tensor_utils.shape(input_ids), segment_id),
        input_positions=tf.fill(tensor_utils.shape(input_ids), i),
        attention_mask=attention_mask,
        input_cache=input_cache,
        reuse=reuse)

    # [batch_size, 1, vocab_size]
    logits = model.compute_logits(sequence_output, reuse=reuse)

    def update_values(old_values, current_value):
      """Update stored values with this time step."""
      shape = [1] * len(tensor_utils.shape(old_values))
      shape[:2] = [batch_size, num_steps]
      tile = tensor_utils.shape(old_values)
      tile[:2] = [1, 1]
      condition = tf.tile(tf.reshape(is_written, shape), tile)
      tile = [1] * len(tensor_utils.shape(old_values))
      tile[1] = num_steps
      current_value = tf.tile(current_value, tile)
      return tf.where(condition, old_values, current_value)

    # [batch_size, num_steps, num_layers, num_heads, head_size]
    beam_output_cache = TransformerCache(
        keys=update_values(state.output_cache.keys, step_cache.keys),
        values=update_values(state.output_cache.values, step_cache.values))

    # Return new state.
    state = DecodeState(
        encoder_cache=state.encoder_cache,
        encoder_cache_mask=state.encoder_cache_mask,
        output_cache=beam_output_cache)

    return tf.squeeze(logits, 1), state
Go from ids to logits for next symbol.
language/capwap/utils/transformer_utils.py