code: 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
  
file_path: tensor2tensor/tensor2tensor/layers/message_passing_attention.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.utils import expert_utils
  import tensorflow.compat.v1
label_desc: Use image node_id 9 for calling a global function with example usage: dot_product_mpnn_attention(q, k, v, adjacency_matrix, num_edge_types, num_transforms, use_weighted_sum, name) and returns: output