def transformer_pointer_prediction_layer(feature_name,
                                         encoder_output,
                                         x,
                                         encoder_decoder_attention_bias,
                                         hparams,
                                         features,
                                         loss_mask,
                                         layer_collection=None):
  """Layer that predicts the start or end token position.

  Args:
    feature_name: 'targets_start_token' or 'targets_end_token'
    encoder_output: [batch_size, input_length, hidden_size] tensor with encoder
      outputs
    x: [batch_size, target_length, 1, hidden_size] tensor with decoder outputs
    encoder_decoder_attention_bias: [batch_size, input_length, target_length]
      attention mask
    hparams: Hyper parameters
    features: Feature dictionary
    loss_mask: [batch_size, target_length] mask for loss computation.
    layer_collection: Layer collection

  Returns:
    (x, logits, loss)
  """
  if isinstance(encoder_output, list):
    pointer_encoder_output = encoder_output[1]
    encoder_output = sum(encoder_output)
  else:
    pointer_encoder_output = encoder_output
  with tf.variable_scope("%s_prediction" % feature_name):
    x = maybe_flatten4d3d(x)
    encoder_decoder_attention_bias = common_layers.flatten4d3d(
        encoder_decoder_attention_bias)
    q = common_attention.compute_attention_component(x, hparams.hidden_size)
    k = common_attention.compute_attention_component(encoder_output,
                                                     hparams.hidden_size)
    # Scaled dot-product attention
    scalar = tf.rsqrt(tf.to_float(common_layers.shape_list(q)[2]))
    logits = tf.matmul(q * scalar, k, transpose_b=True)

    logits += encoder_decoder_attention_bias

    labels = features["%s_raw" % feature_name]
    xent = tf.nn.sparse_softmax_cross_entropy_with_logits(
        logits=logits, labels=labels)
    loss = tf.reduce_sum(xent * loss_mask)

    pointer_out = gather_2d(pointer_encoder_output, labels)
    y = common_layers.layer_preprocess(
        pointer_out, hparams, layer_collection=layer_collection)
    x = common_layers.layer_postprocess(x, y, hparams)
    return x, logits, loss
Layer that predicts the start or end token position.

  Args:
    feature_name: 'targets_start_token' or 'targets_end_token'
    encoder_output: [batch_size, input_length, hidden_size] tensor with encoder
      outputs
    x: [batch_size, target_length, 1, hidden_size] tensor with decoder outputs
    encoder_decoder_attention_bias: [batch_size, input_length, target_length]
      attention mask
    hparams: Hyper parameters
    features: Feature dictionary
    loss_mask: [batch_size, target_length] mask for loss computation.
    layer_collection: Layer collection

  Returns:
    (x, logits, loss)
  
tensor2tensor/models/research/transformer_seq2edits.py