Dot product attention with edge vectors.

  Let B be the number of batches.
  Let N be the number of nodes in the graph.
  Let K be the size of the attention keys/queries.
  Let V be the size of the attention values.
  Let T be the total number of transforms (num_transforms).

  Args:
    q: The query Tensor of shape [B, N, K].
    k: The key Tensor of shape [B, T, N, K].
    v: The value Tensor of shape [B, T, N, V].
    adjacency_matrix: A Tensor of shape [B, N, N, T]. An entry at
      indices b, i, j, k is the indicator of the edge
      from node j to node i in batch b. A standard adjacency matrix will only
      have one edge type while a mutigraph will have multiple edge types.
    num_edge_types: An integer specifying number of edge types.
    num_transforms: An integer indicating number of transforms (T). If None,
      then num_transforms will be equal to num_edge_types.
    use_weighted_sum: If False, will only use a single transform per edge type.
      Otherwise, use a learned weighted sum of transforms per edge type.
    name: A string.

  Returns:
    A Tensor of shape [B, N, V] storing the result of computing attention
    weights using the queries and keys and combining the values according to
    those weights.

  Raises:
    ValueError: if num_transforms doesn't equal num_edge_types and not using
      weighted sum.

def dot_product_mpnn_attention(q,
                               k,
                               v,
                               adjacency_matrix,
                               num_edge_types,
                               num_transforms=None,
                               use_weighted_sum=False,
                               name=None):
  """Dot product attention with edge vectors.

  Let B be the number of batches.
  Let N be the number of nodes in the graph.
  Let K be the size of the attention keys/queries.
  Let V be the size of the attention values.
  Let T be the total number of transforms (num_transforms).

  Args:
    q: The query Tensor of shape [B, N, K].
    k: The key Tensor of shape [B, T, N, K].
    v: The value Tensor of shape [B, T, N, V].
    adjacency_matrix: A Tensor of shape [B, N, N, T]. An entry at
      indices b, i, j, k is the indicator of the edge
      from node j to node i in batch b. A standard adjacency matrix will only
      have one edge type while a mutigraph will have multiple edge types.
    num_edge_types: An integer specifying number of edge types.
    num_transforms: An integer indicating number of transforms (T). If None,
      then num_transforms will be equal to num_edge_types.
    use_weighted_sum: If False, will only use a single transform per edge type.
      Otherwise, use a learned weighted sum of transforms per edge type.
    name: A string.

  Returns:
    A Tensor of shape [B, N, V] storing the result of computing attention
    weights using the queries and keys and combining the values according to
    those weights.

  Raises:
    ValueError: if num_transforms doesn't equal num_edge_types and not using
      weighted sum.
  """
  with tf.variable_scope(
      name,
      default_name="dot_product_mpnn_attention",
      values=[q, k, v, adjacency_matrix, num_edge_types]):
    # If not explicitly set, use num_transforms set to num_edge_types.
    num_transforms = (
        num_edge_types if num_transforms is None else num_transforms)

    if not use_weighted_sum and num_transforms != num_edge_types:
      raise ValueError("num_transforms must equal num_edge_types unless "
                       "use_weighted_sum is True")

    # Computes the raw dot-product attention values between each query and
    # the corresponding keys it needs to consider.
    #
    # This operation takes the dot product of (the query for
    # each node) and (the key for each node for each possible edge type),
    # creating an N x N matrix for each edge type. The entry at index (i, j)
    # is the dot-product for the edge from node i to node j of the appropriate
    # type. These dot products will eventually become attention weights
    # specifying how much node i weights an edge of that type coming from node
    # j.
    all_edge_logits = tf.matmul(
        tf.tile(tf.expand_dims(q, axis=1), [1, num_edge_types, 1, 1]),
        k,
        transpose_b=True)

    # The adjacency matrix assumes there is only one directed edge (i <- j) for
    # each pair of nodes. If such an edge exists, it contains the integer
    # type of that edge at position (i, j) of the adjacency matrix.
    #
    # Construct edge_vectors of shape [B, N, N, T].
    if use_weighted_sum:
      # Use dense representation for edge vectors.
      edge_vectors = make_edge_vectors(
          adjacency_matrix,
          num_edge_types,
          num_transforms)
    else:
      # Generate one-hot vectors based on edge types.
      # If there is an edge from node j to node i of type t, then index t of the
      # last dimension is 1 for entry (i, j) of the second and third dimensions.
      edge_vectors = tf.one_hot(adjacency_matrix, num_transforms)

    # Rearranging the dimensions to match the shape of all_edge_logits.
    edge_vectors = tf.transpose(edge_vectors, [0, 3, 1, 2])

    # Element-wise multiplies all_edge_logits and edge_vectors.
    #
    # In other words: all_edge_logits contains N x N matrices of query-key
    # products. This element-wise multiplication zeroes out entries that do not
    # correspond to actual edges in the graph of the appropriate edge type.
    # all_edge_logits retains shape [B, T, N, N].
    all_edge_logits *= edge_vectors

    # Since there can only be one edge from node A to node B, we can collapse
    # the T different adjacency matrices containing key-query pairs into one
    # adjacency matrix. logits is [B, N, N].
    # TODO(dbieber): Use a reshape instead of reduce sum to attend over all
    # edges instead of over all neighboring nodes to handle the multigraph case.
    logits = tf.reduce_sum(all_edge_logits, axis=1)

    # For pairs of nodes with no edges between them, add a large negative bias
    # to each location without an edge so that the softmax of entries with the
    # value 0 become a small negative number instead.
    bias = 0
    bias = tf.to_float(tf.equal(
        tf.reduce_sum(adjacency_matrix, axis=-1), 0)) * -1e9
    logits += bias

    # Turn the raw key-query products into a probability distribution (or,
    # in terms of attention, weights). The softmax is computed across the
    # last dimension of logits.
    compatibility = tf.nn.softmax(logits)  # Shape [B, N, N].

    # Computes a summary showing the attention matrix as an image. Does not do
    # any work toward actually performing attention.
    common_attention.attention_image_summary(
        tf.expand_dims(compatibility, axis=1), None)

    # Repeats the attention matrix T times for each batch, producing
    # a tensor with shape [B, T, N, N] where the [N, N] component is T
    # repeats of the values found in compatibility.
    edge_compatibility = tf.tile(
        tf.expand_dims(compatibility, axis=1), [1, num_edge_types, 1, 1])

    # Zeroes out the entries in edge_compatibility that do not correspond to
    # actual edges.
    edge_compatibility *= edge_vectors  # Shape [B, T, N, N].

    output = compute_values(edge_compatibility, v)
    return output