code: def dense_message_pass(node_states, edge_matrices):
    
    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
docs: 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.