def attention(query_antecedent, memory_antecedent, bias, model_size, num_heads):
  """Multihead dot-product attention.

  Args:
    query_antecedent: <float>[batch_size, query_length, query_emb_size]
    memory_antecedent: <float>[batch_size, memory_length, memory_emb_size]
    bias: <float>[?, ?, ?, ?] tensor that can be broadcast to have shape
      [batch_size, num_heads, query_length, memory_length].
    model_size: Integer specifying number of dimensions.
    num_heads: Integer specifying number of attention heads.

  Returns:
    <float>[batch_size, length, model_size]
  """
  q, k, v = _prepare_attention_inputs(query_antecedent, memory_antecedent,
                                      model_size, model_size, num_heads)
  # q is <float>[batch_size, num_heads, query_length, model_size / num_heads]
  # k, v <float>[batch_size, num_heads, memory_length, model_size / num_heads]
  logits = tf.matmul(q, k, transpose_b=True)
  # logits is <float>[batch_size, num_heads, query_length, memory_length]
  logits += bias
  weights = tf.nn.softmax(logits, name="attention_weights")
  x = tf.matmul(weights, v)
  # Combine heads and perform output transformation.
  x = _combine_heads(x)
  x = linear_transform(x, model_size, "output_transform")
  return x
Multihead dot-product attention.

  Args:
    query_antecedent: <float>[batch_size, query_length, query_emb_size]
    memory_antecedent: <float>[batch_size, memory_length, memory_emb_size]
    bias: <float>[?, ?, ?, ?] tensor that can be broadcast to have shape
      [batch_size, num_heads, query_length, memory_length].
    model_size: Integer specifying number of dimensions.
    num_heads: Integer specifying number of attention heads.

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