def compute_joint_mlp_logits(sequence, max_span_length):
  """Computes joint span (start, end) logits from sequence input."""
  batch_size, seq_length, hidden_size = modeling.get_shape_list(
      sequence, expected_rank=3)

  projection_size = hidden_size  # This seems to be a reasonable setting.

  with tf.variable_scope("joint_span"):
    projection = tf.layers.dense(
        sequence,
        projection_size * 2,
        activation=None,
        kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),
        name="projection")
    start_projection, end_projection = tf.split(projection, 2, axis=-1)

    # 1. The start representations are tiled max_answer_length times.
    # TODO(danielandor): Use the mask to compute an optimal span list.
    starts = tf.reshape(start_projection,
                        [batch_size * seq_length, 1, projection_size])
    starts = tf.tile(starts, [1, max_span_length, 1])
    starts = tf.reshape(
        starts, [batch_size, seq_length * max_span_length, projection_size])

    # 2. To make the end representations, we compute band diagonal indices and
    #    perform a batched gather.
    seqs = tf.expand_dims(tf.range(seq_length), 1)
    offsets = tf.expand_dims(tf.range(max_span_length), 0)
    indices = seqs + offsets  # uses broadcasting
    indices.shape.assert_is_compatible_with((seq_length, max_span_length))
    indices = tf.reshape(indices, [1, seq_length * max_span_length])
    indices = tf.tile(indices, [batch_size, 1])
    indices = tf.minimum(indices, seq_length - 1)  # clips indices
    ends = tf.batch_gather(end_projection, indices)

    # 3. The final step adds the starts and ends.
    ends.shape.assert_is_compatible_with(starts.shape)
    inputs = starts + ends
    inputs = modeling.gelu(inputs)  # Bias is already in the projection.
    inputs = contrib_layers.layer_norm(inputs)
    start_logits = tf.layers.dense(
        inputs,
        1,
        activation=None,
        kernel_initializer=tf.truncated_normal_initializer(stddev=0.02),
        name="logits")
  return tf.reshape(start_logits, [batch_size, seq_length, max_span_length])
Computes joint span (start, end) logits from sequence input.
language/capwap/synthetic/generate_answers.py