def _get_action_logits(encoder_output,
                       decoder_output,
                       output_vocab_embeddings_table,
                       output_vocab_size,
                       model_config,
                       input_copy_mask=None,
                       clean_output_mask=None,
                       use_gating_mechanism=True):
  """Generate output logits given decoder output.

  This effectively combines a Pointer Network (Vinyals et al., 2015) with a
  standard softmax output layer for selecting symbols from an output vocabulary,
  similar to:
      - Jia and Liang, 2016 (https://arxiv.org/abs/1606.03622)
      - Gulcehre et al., 2016 (https://arxiv.org/abs/1603.08148)
      - Gu et al., 2016 (https://arxiv.org/abs/1603.06393)
      - See et al. 2017 (https://arxiv.org/abs/1704.04368)

  Args:
    encoder_output: Tensor representing encoder output of shape (batch size,
      input length, encoder dims).
    decoder_output: Tensor representing decoder output of shape (batch size, #
      decoded steps, decoder dims).
    output_vocab_embeddings_table: Embeddings for output vocabulary of shape
      (output_vocab_size, target embedding dims).
    output_vocab_size: Integer size of output_vocab_embeddings_table outer dim.
    model_config: ModelConfig proto.
    input_copy_mask: Mask of the input sequence for copying.
    clean_output_mask: Mask of the output vocab. For clean
      inference only.
    use_gating_mechanism: Whether to use gating mechanism.

  Returns:
    Tensor of shape (batch_size, output_vocab_size + input length) representing
    unnormalized logits for both copy and generate actions.
  """

  with tf.variable_scope("logits_transforms"):
    decoder_dims = decoder_output.get_shape()[-1]
    target_embedding_dims = model_config.model_parameters.target_embedding_dims

    # Dot product the decoder output with representations of each of the output
    # symbols to get a set of unnormalized logits for each output vocab item.
    # We need to tile the output vocab embeddings across the batch.
    output_vocab_transform = tf.expand_dims(output_vocab_embeddings_table, 0)
    batch_size = tf.shape(decoder_output)[0]
    output_vocab_transform = tf.tile(output_vocab_transform, [batch_size, 1, 1])
    # Transform representations to the target_embedding_dims.
    if decoder_dims != target_embedding_dims:
      transformed_decoder_output = common_layers.linear_transform(
          decoder_output, target_embedding_dims, "decoder_transform")
    else:
      transformed_decoder_output = decoder_output
    generate_logits = tf.matmul(
        transformed_decoder_output, output_vocab_transform, transpose_b=True)
    generate_logits_bias = tf.get_variable(
        "generate_logits_bias", shape=(output_vocab_size))
    generate_logits += generate_logits_bias

    # Dot product the decoder output with representations from the encoder
    # output.
    # This is necessary vs. re-using the encoder-decoder attention weights
    # because those use multihead attention.
    # First, need to transform representations to the decoder dimensions.
    transformed_encoder_output = common_layers.linear_transform(
        encoder_output, decoder_dims, "encoder_transform")

    copy_logits = tf.matmul(
        decoder_output, transformed_encoder_output, transpose_b=True)
    # This contains scores representing the probability of copying from input
    # (3rd dim) to output (2nd dim).

    # Optionally apply a soft gating mechanism to determine whether
    # to select from copy or generate logits.
    # TODO(petershaw): Evaluate and improve this gating mechanism.
    # The current implementation is most likely not optimal, since it applies
    # a scalar in the range [0,1] prior to softmax.
    if use_gating_mechanism:
      prob_gen_unnormalized = common_layers.linear_transform(
          decoder_output, 1, "prob_gen")
      prob_gen_bias = tf.get_variable("prob_gen_bias", shape=(1))
      prob_gen_unnormalized += prob_gen_bias
      prob_gen = tf.sigmoid(prob_gen_unnormalized)
      # Squeeze so that prob_gen has shape [batch_size, decode_length]
      prob_gen = tf.squeeze(prob_gen, axis=2)

      # These are the 'generate' logits so are scaled by P_gen.
      generate_logits *= tf.expand_dims(prob_gen, axis=-1)
      # These are the 'copy' logits so are scaled by 1 - P_gen.
      copy_logits *= tf.expand_dims(1 - prob_gen, axis=-1)

    if clean_output_mask is not None:
      clean_mask = (1 - tf.dtypes.cast(
          clean_output_mask, dtype=tf.dtypes.float32)) * LOGIT_MASK_VALUE

      batch_size = common_layers.get_shape_list(generate_logits)[0]
      output_vocab_size = common_layers.get_shape_list(generate_logits)[-1]

      clean_mask = tf.reshape(
          tf.tile(clean_mask, [batch_size]), [batch_size, output_vocab_size])
      generate_logits += tf.expand_dims(clean_mask, axis=1)

    if input_copy_mask is not None:
      copy_mask = (1 - tf.dtypes.cast(
          input_copy_mask, dtype=tf.dtypes.float32)) * LOGIT_MASK_VALUE
      copy_logits += tf.expand_dims(copy_mask, axis=1)

    # Concatenate logits into a single vector; first N (fixed) inputs are the
    # generation probabilities, and next are the copy probabilities for each
    # input (well, they aren't really probabilities, but scores.)
    extended_logits = tf.concat([generate_logits, copy_logits], axis=2)
    return extended_logits
Generate output logits given decoder output.

  This effectively combines a Pointer Network (Vinyals et al., 2015) with a
  standard softmax output layer for selecting symbols from an output vocabulary,
  similar to:
      - Jia and Liang, 2016 (https://arxiv.org/abs/1606.03622)
      - Gulcehre et al., 2016 (https://arxiv.org/abs/1603.08148)
      - Gu et al., 2016 (https://arxiv.org/abs/1603.06393)
      - See et al. 2017 (https://arxiv.org/abs/1704.04368)

  Args:
    encoder_output: Tensor representing encoder output of shape (batch size,
      input length, encoder dims).
    decoder_output: Tensor representing decoder output of shape (batch size, #
      decoded steps, decoder dims).
    output_vocab_embeddings_table: Embeddings for output vocabulary of shape
      (output_vocab_size, target embedding dims).
    output_vocab_size: Integer size of output_vocab_embeddings_table outer dim.
    model_config: ModelConfig proto.
    input_copy_mask: Mask of the input sequence for copying.
    clean_output_mask: Mask of the output vocab. For clean
      inference only.
    use_gating_mechanism: Whether to use gating mechanism.

  Returns:
    Tensor of shape (batch_size, output_vocab_size + input length) representing
    unnormalized logits for both copy and generate actions.
  
language/xsp/model/transformer.py