Precompute the a_in and a_out tensors.

  (we don't want to add to the graph everytime _fprop is called)
  Args:
    adjacency: placeholder of real valued vectors of shape [B, L, L, E]
    hparams: HParams object
  Returns:
    edge_matrices: [batch, L * D, L * D] the dense matrix for message passing
    viewed as a block matrix (L,L) blocks of size (D,D). Each plot is a function
    of the edge vector of the adjacency matrix at that spot.

def precompute_edge_matrices(adjacency, hparams):
  """Precompute the a_in and a_out tensors.

  (we don't want to add to the graph everytime _fprop is called)
  Args:
    adjacency: placeholder of real valued vectors of shape [B, L, L, E]
    hparams: HParams object
  Returns:
    edge_matrices: [batch, L * D, L * D] the dense matrix for message passing
    viewed as a block matrix (L,L) blocks of size (D,D). Each plot is a function
    of the edge vector of the adjacency matrix at that spot.
  """
  batch_size, num_nodes, _, edge_dim = common_layers.shape_list(adjacency)

  # build the edge_network for incoming edges
  with tf.variable_scope("edge_network"):
    x = tf.reshape(
        adjacency, [batch_size * num_nodes * num_nodes, edge_dim],
        name="adj_reshape_in")

    for ip_layer in range(hparams.edge_network_layers):
      name = "edge_network_layer_%d"%ip_layer
      x = tf.layers.dense(common_layers.layer_preprocess(x, hparams),
                          hparams.edge_network_hidden_size,
                          activation=tf.nn.relu,
                          name=name)
    x = tf.layers.dense(common_layers.layer_preprocess(x, hparams),
                        hparams.hidden_size**2,
                        activation=None,
                        name="edge_network_output")

  # x = [batch * l * l, d *d]
  edge_matrices_flat = tf.reshape(x, [batch_size, num_nodes,
                                      num_nodes, hparams.hidden_size,
                                      hparams.hidden_size])

  # reshape to [batch, l * d, l *d]
  edge_matrices = tf.reshape(
      tf.transpose(edge_matrices_flat, [0, 1, 3, 2, 4]), [
          -1, num_nodes * hparams.hidden_size,
          num_nodes * hparams.hidden_size
      ],
      name="edge_matrices")

  return edge_matrices