def graph_attention(query_antecedent,
                    memory_antecedent,
                    bias,
                    model_size,
                    num_heads,
                    adjacency_matrix,
                    num_labels,
                    use_one_hot_embeddings=False):
  """Multihead attention with edge vectors.

  Args:
    query_antecedent: <float>[batch_size, length, emb_size]
    memory_antecedent: <float>[batch_size, length, emb_size]
    bias: <float>[?, ?, length, length] tensor that can be broadcast to have
      shape [batch_size, num_heads, length, length].
    model_size: Integer specifying number of dimensions.
    num_heads: Integer specifying number of attention heads.
    adjacency_matrix: <int>[batch_size, length, length] specifying edge labels.
    num_labels: Integer specifying number of unique edge labels.
    use_one_hot_embeddings: bool. If True, use one-hot method for word
      embeddings. If False, use `tf.nn.embedding_lookup()`.

  Returns:
    <float>[batch_size, length, model_size]
  """
  q, k, v = _prepare_attention_inputs(query_antecedent, memory_antecedent,
                                      model_size, model_size, num_heads)
  logits = _edge_vectors_attention_inner(
      q,
      k,
      adjacency_matrix,
      num_labels,
      transpose=True,
      name="keys",
      use_one_hot_embeddings=use_one_hot_embeddings)
  logits += bias
  weights = tf.nn.softmax(logits, name="attention_weights")
  x = _edge_vectors_attention_inner(
      weights,
      v,
      adjacency_matrix,
      num_labels,
      transpose=False,
      name="values",
      use_one_hot_embeddings=use_one_hot_embeddings)
  # Combine heads and perform output transformation.
  x = _combine_heads(x)
  x = linear_transform(x, model_size, "output_transform")
  return x
Multihead attention with edge vectors.

  Args:
    query_antecedent: <float>[batch_size, length, emb_size]
    memory_antecedent: <float>[batch_size, length, emb_size]
    bias: <float>[?, ?, length, length] tensor that can be broadcast to have
      shape [batch_size, num_heads, length, length].
    model_size: Integer specifying number of dimensions.
    num_heads: Integer specifying number of attention heads.
    adjacency_matrix: <int>[batch_size, length, length] specifying edge labels.
    num_labels: Integer specifying number of unique edge labels.
    use_one_hot_embeddings: bool. If True, use one-hot method for word
      embeddings. If False, use `tf.nn.embedding_lookup()`.

  Returns:
    <float>[batch_size, length, model_size]
  
language/xsp/model/common_layers.py