code: def dense_message_pass(node_states, edge_matrices):
      """Computes a_t from h_{t-1}, see bottom of page 3 in the paper.
  
      Args:
        node_states: [B, L, D] tensor (h_{t-1})
        edge_matrices (tf.float32): [B, L*D, L*D]
  
      Returns:
        messages (tf.float32): [B, L, D] For each pair
          of nodes in the graph a message is sent along both the incoming and
          outgoing edge.
      """
      batch_size, num_nodes, node_dim = common_layers.shape_list(
          node_states
      )
  
      # Stack the nodes as a big column vector.
      h_flat = tf.reshape(
          node_states,
          [batch_size, num_nodes * node_dim, 1],
          name="h_flat",
      )
  
      messages = tf.reshape(
          tf.matmul(edge_matrices, h_flat),
          [batch_size * num_nodes, node_dim],
          name="messages_matmul",
      )
  
      message_bias = tf.get_variable("message_bias", shape=node_dim)
      messages = messages + message_bias
      messages = tf.reshape(messages, [batch_size, num_nodes, node_dim])
      return messages
  
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 13 for calling a global function with example usage: dense_message_pass(node_states, edge_matrices) and returns: messages