code: def ggnn_fast_dense(
      node_states,
      adjacency_matrix,
      num_edge_types,
      total_value_depth,
      name=None,
  ):
      """ggnn version of the MPNN from Gilmer et al.
  
      Let B be the number of batches.
      Let D be the size of the node hidden states.
      Let K be the size of the attention keys/queries.
      Let V be the size of the output of the ggnn.
      Let T be the number of transforms / edge types.
  
      Args:
        node_states: The value Tensor of shape [B, T, N, D].
        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 values of one, while a
          mutigraph may have larger integer values.
        num_edge_types: An integer specifying number of edge types.
        total_value_depth: An integer (V)
        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.
      """
      # between the same nodes (with only one edge of each type. adjacency_matrix
      # will need to be converted to shape [B, T, N, N].
      with tf.variable_scope(
          name,
          default_name="ggnn_fast_dense",
          values=[node_states, adjacency_matrix, num_edge_types],
      ):
          nodes_shape = common_layers.shape_list(node_states)
          v = _compute_edge_transforms(
              node_states,
              total_value_depth,
              num_edge_types,
              name="v_mpnn",
          )
          v = tf.reshape(
              v,
              [
                  nodes_shape[0],
                  nodes_shape[1],
                  num_edge_types,
                  total_value_depth,
              ],
          )  # Shape [B, N, T, V].
          v = tf.transpose(v, [0, 2, 1, 3])  # Shape [B, T, N, V].
  
          # Rearranging the dimensions to match the shape of all_edge_logits.
          edge_vectors = tf.transpose(adjacency_matrix, [0, 3, 1, 2])
          output = compute_values(edge_vectors, 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 10 for calling a global function with example usage: ggnn_fast_dense(node_states, adjacency_matrix, num_edge_types, total_value_depth, name) and returns: output