def _get_target_embeddings(encoder_input, output_vocab_embeddings_table,
                           decode_steps, model_config):
  """Get target embeddings from either output table or copied from input.

  Args:
    encoder_input: Tensor representing encoder output of shape (batch size,
      input length, encoder dims).
    output_vocab_embeddings_table: Embeddings for output vocabulary of shape
      (output_vocab_size, target embedding dims).
    decode_steps: DecodeSteps tuple with tensors of shape (batch size, # steps).
    model_config: ModelConfig proto.

  Returns:
    Tensor of shape (batch_size, # steps, target embedding dims) representing
    unnormalized logits for both copy and generate actions.
  """
  input_length = tf.shape(encoder_input)[1]
  # Size of one_hot is (batch size, # steps, input length).
  one_hot = tf.one_hot(decode_steps.action_ids, input_length)
  # Size of encoder_dims is (batch size, input length, encoder dims).
  # Size of matrix multiplication is then (batch size, # steps, encoder_dims).
  copy_embeddings = tf.matmul(one_hot, encoder_input)

  # Need a linear transformation to ensure copy embeddings are right size.
  # Shape will then be (batch size, # steps, target_embedding_dims)
  copy_embeddings = common_layers.linear_transform(
      copy_embeddings, model_config.model_parameters.target_embedding_dims,
      "copy_embeddings_transform")
  # Simply get the generate embeddings from the output vocab table.

  generate_steps = tf.equal(
      decode_steps.action_types,
      tf.constant(constants.GENERATE_ACTION, dtype=tf.int64))
  generate_embeddings = common_layers.embedding_lookup(
      output_vocab_embeddings_table,
      decode_steps.action_ids * tf.to_int64(generate_steps))
  # For a given step, only use either copy OR generate embeddings.
  copy_steps = tf.equal(decode_steps.action_types,
                        tf.constant(constants.COPY_ACTION, dtype=tf.int64))

  copy_mask = tf.to_float(tf.expand_dims(copy_steps, axis=-1))
  generate_mask = tf.to_float(tf.expand_dims(generate_steps, axis=-1))
  target_embeddings = copy_embeddings * copy_mask + generate_embeddings * generate_mask

  return target_embeddings
Get target embeddings from either output table or copied from input.

  Args:
    encoder_input: Tensor representing encoder output of shape (batch size,
      input length, encoder dims).
    output_vocab_embeddings_table: Embeddings for output vocabulary of shape
      (output_vocab_size, target embedding dims).
    decode_steps: DecodeSteps tuple with tensors of shape (batch size, # steps).
    model_config: ModelConfig proto.

  Returns:
    Tensor of shape (batch_size, # steps, target embedding dims) representing
    unnormalized logits for both copy and generate actions.
  
language/xsp/model/transformer.py