def _prepare_attention_inputs(query_antecedent, memory_antecedent, query_size,
                              memory_size, num_heads):
  """Prepare multi-head attention inputs.

  Args:
    query_antecedent: <float>[batch_size, query_length, query_input_size]
    memory_antecedent: <float>[batch_size, memory_length, memory_input_size]
    query_size: Integer specifying number of dimensions for q.
    memory_size: Integer specifying number of deimsnions for k and v.
    num_heads: Integer specifying number of attention heads.

  Returns:
    q: <float>[batch_size, num_heads, query_length, query_size / num_heads]
    k: <float>[batch_size, num_heads, memory_length, memory_size / num_heads]
    v: <float>[batch_size, num_heads, memory_length, memory_size / num_heads]
  """
  q = linear_transform(query_antecedent, query_size, "q_transform")
  k = linear_transform(memory_antecedent, memory_size, "k_transform")
  v = linear_transform(memory_antecedent, memory_size, "v_transform")
  q = _split_heads(q, num_heads)
  k = _split_heads(k, num_heads)
  v = _split_heads(v, num_heads)
  # Scale query inputs.
  key_depth_per_head = query_size // num_heads
  q *= key_depth_per_head**-0.5
  return q, k, v
Prepare multi-head attention inputs.

  Args:
    query_antecedent: <float>[batch_size, query_length, query_input_size]
    memory_antecedent: <float>[batch_size, memory_length, memory_input_size]
    query_size: Integer specifying number of dimensions for q.
    memory_size: Integer specifying number of deimsnions for k and v.
    num_heads: Integer specifying number of attention heads.

  Returns:
    q: <float>[batch_size, num_heads, query_length, query_size / num_heads]
    k: <float>[batch_size, num_heads, memory_length, memory_size / num_heads]
    v: <float>[batch_size, num_heads, memory_length, memory_size / num_heads]
  
language/xsp/model/common_layers.py