def symbols_to_logits_fn(current_index, logit_indices):
    """Go from targets to logits.

    Args:
      current_index: Integer corresponding to 0-indexed decoder step.
      logit_indices: Tensor of shape [batch_size * beam_width, decode_length +
        1] to input to decoder.

    Returns:
      Tensor of shape [batch_size * beam_width, output_vocab_size] representing
      logits for the current decoder step.

    Raises:
      ValueError if inputs do not have static length.
    """
    decode_steps = decode_utils.get_decode_steps(logit_indices,
                                                 output_vocab_size,
                                                 model_config)
    target_embeddings = _get_target_embeddings(input_embeddings,
                                               output_vocab_embeddings_table,
                                               decode_steps, model_config)
    decoder_output = _build_transformer_decoder(
        encoder_output,
        source_len,
        target_embeddings,
        mode,
        model_config,
        single_step_index=current_index)
    logits = _get_action_logits(
        encoder_output,
        decoder_output,
        output_vocab_embeddings_table,
        output_vocab_size,
        model_config,
        input_copy_mask=input_copy_mask,
        clean_output_mask=clean_output_mask)
    # Squeeze length dimension, as it should be 1.
    logits = tf.squeeze(logits, axis=[1])
    # Shape of logits should now be:
    # [batch_size * beam_width, output_vocab_size].
    return logits
Go from targets to logits.

    Args:
      current_index: Integer corresponding to 0-indexed decoder step.
      logit_indices: Tensor of shape [batch_size * beam_width, decode_length +
        1] to input to decoder.

    Returns:
      Tensor of shape [batch_size * beam_width, output_vocab_size] representing
      logits for the current decoder step.

    Raises:
      ValueError if inputs do not have static length.
    
language/xsp/model/transformer.py