code: def graph_attention(
      q,
      k,
      v,
      bias,
      dropout_rate=0.0,
      image_shapes=None,
      name=None,
      make_image_summary=True,
      save_weights_to=None,
      dropout_broadcast_dims=None,
      adjacency_matrix=None,
      num_edge_types=5,
  ):
      """graph attention.
  
      Args:
        q: a Tensor with shape [batch, heads, length_q, depth_k]
        k: a Tensor with shape [batch, heads, length_kv, depth_k]
        v: a Tensor with shape [batch, heads, length_kv, depth_v]
        bias: bias Tensor (see attention_bias())
        dropout_rate: a floating point number
        image_shapes: optional tuple of integer scalars.
          see comments for attention_image_summary()
        name: an optional string
        make_image_summary: True if you want an image summary.
        save_weights_to: an optional dictionary to capture attention weights
          for vizualization; the weights tensor will be appended there under
          a string key created from the variable scope (including name).
        dropout_broadcast_dims:  an optional list of integers less than 4
          specifying in which dimensions to broadcast the dropout decisions.
          saves memory.
        adjacency_matrix: optional matrix of [batch, length, length] ids indicating
          edge type
        num_edge_types: an int indicating number of edge types
      Returns:
        A Tensor of shape [batch, length, depth(q)]
      """
      with tf.variable_scope(
          name, default_name="dot_product_attention", values=[q, k, v]
      ) as scope:
          # [batch, num_heads, query_length, memory_length]
          logits = tf.matmul(q, k, transpose_b=True)
          if adjacency_matrix is not None:
              key_head_depth = common_layers.shape_list(q)[-1]
              adjacency_vectors = make_edge_vectors(
                  adjacency_matrix,
                  num_edge_types,
                  key_head_depth,
                  name=name,
              )
              # transposing q to be [batch, length_q, heads, depth_k]
              # to allow for matmul with [batch, length_q, length_q, depth_k]
              q_t = tf.transpose(q, [0, 2, 1, 3])
              adj_logits = tf.matmul(
                  q_t, adjacency_vectors, transpose_b=True
              )
              logits += tf.transpose(adj_logits, [0, 2, 1, 3])
              # [batch, depth, num_nodes, num_nodes]
          if bias is not None:
              logits += bias
          weights = tf.nn.softmax(logits, name="attention_weights")
          if save_weights_to is not None:
              save_weights_to[scope.name] = weights
          # dropping out the attention links for each of the heads
          weights = common_layers.dropout_with_broadcast_dims(
              weights,
              1.0 - dropout_rate,
              broadcast_dims=dropout_broadcast_dims,
          )
          if (
              common_layers.should_generate_summaries()
              and make_image_summary
          ):
              common_attention.attention_image_summary(
                  weights, image_shapes
              )
          return tf.matmul(weights, v)
  
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 3 for calling a global function with example usage: graph_attention(q, k, v, bias, dropout_rate, image_shapes, name, make_image_summary, save_weights_to, dropout_broadcast_dims, adjacency_matrix, num_edge_types) and returns: tf