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.

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