code: def multihead_graph_attention(
      query_antecedent,
      memory_antecedent,
      bias,
      total_key_depth,
      total_value_depth,
      output_depth,
      num_heads,
      dropout_rate,
      image_shapes=None,
      attention_type="edge_vector",
      name="multihead_graph_attention",
      save_weights_to=None,
      make_image_summary=True,
      dropout_broadcast_dims=None,
      adjacency_matrix=None,
      num_edge_types=5,
      vars_3d=False,
      **kwargs,
  ):
      """Multihead scaled-dot-product attention with input/output transformations.
  
      Args:
        query_antecedent: a Tensor with shape [batch, length_q, channels]
        memory_antecedent: a Tensor with shape [batch, length_m, channels] or None
        bias: bias Tensor (see attention_bias())
        total_key_depth: an integer
        total_value_depth: an integer
        output_depth: an integer
        num_heads: an integer dividing total_key_depth and total_value_depth
        dropout_rate: a floating point number
        image_shapes: optional tuple of integer scalars.
                      see comments for attention_image_summary()
        attention_type: a string, either "dot_product", "dot_product_relative",
                        "local_mask_right", "local_unmasked", "masked_dilated_1d",
                        "unmasked_dilated_1d", graph, or any attention function
                        with the signature (query, key, value, **kwargs)
        name: an optional string.
        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).
        make_image_summary: Whether to make an attention image summary.
        dropout_broadcast_dims:  an optional list of integers less than 4
          specifying in which dimensions to broadcast the dropout decisions.
          saves memory.
        adjacency_matrix: an optional tensor of shape [batch, len_q, len_q]
          containing edge vectors for attention
        num_edge_types: number of edge types, an int
        vars_3d: use 3-dimensional variables for input/output transformations
        **kwargs (dict): Parameters for the attention function
  
      Returns:
        The result of the attention transformation. The output shape is
            [batch_size, length_q, output_depth]
  
      Raises:
        ValueError: if the key depth or value depth are not divisible by the
          number of attention heads.
      """
      if total_key_depth % num_heads != 0:
          raise ValueError(
              "Key depth (%d) must be divisible by the number of "
              "attention heads (%d)." % (total_key_depth, num_heads)
          )
      if total_value_depth % num_heads != 0:
          raise ValueError(
              "Value depth (%d) must be divisible by the number of "
              "attention heads (%d)." % (total_value_depth, num_heads)
          )
      vars_3d_num_heads = num_heads if vars_3d else None
      with tf.variable_scope(
          name,
          default_name="multihead_attention",
          values=[query_antecedent, memory_antecedent],
      ):
          q, k, v = common_attention.compute_qkv(
              query_antecedent,
              memory_antecedent,
              total_key_depth,
              total_value_depth,
              vars_3d_num_heads=vars_3d_num_heads,
          )
          q = common_attention.split_heads(q, num_heads)
          k = common_attention.split_heads(k, num_heads)
          v = common_attention.split_heads(v, num_heads)
  
          key_depth_per_head = total_key_depth // num_heads
          if not vars_3d:
              q *= key_depth_per_head**-0.5
  
          additional_returned_value = None
          if callable(
              attention_type
          ):  # Generic way to extend multihead_attention
              x = attention_type(q, k, v, **kwargs)
              if isinstance(x, tuple):
                  x, additional_returned_value = x  # Unpack
  
          elif attention_type == "edge_vector":
              x = graph_attention(
                  q,
                  k,
                  v,
                  bias,
                  dropout_rate,
                  image_shapes,
                  save_weights_to=save_weights_to,
                  make_image_summary=make_image_summary,
                  dropout_broadcast_dims=dropout_broadcast_dims,
                  adjacency_matrix=adjacency_matrix,
                  num_edge_types=num_edge_types,
              )
  
          x = common_attention.combine_heads(x)
  
          # Set last dim specifically.
          x.set_shape(x.shape.as_list()[:-1] + [total_value_depth])
  
          if vars_3d:
              o_var = tf.get_variable(
                  "o",
                  [
                      num_heads,
                      total_value_depth // num_heads,
                      output_depth,
                  ],
              )
              o_var = tf.reshape(
                  o_var, [total_value_depth, output_depth]
              )
              x = tf.tensordot(x, o_var, axes=1)
          else:
              x = common_layers.dense(
                  x,
                  output_depth,
                  use_bias=False,
                  name="output_transform",
              )
          if additional_returned_value is not None:
              return x, additional_returned_value
          return x
  
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 1 for calling a global function with example usage: multihead_graph_attention(query_antecedent, memory_antecedent, bias, total_key_depth, total_value_depth, output_depth, num_heads, dropout_rate, image_shapes, attention_type, name, save_weights_to, make_image_summary, dropout_broadcast_dims, adjacency_matrix, num_edge_types, vars_3d) and returns: x, x,  additional_returned_value