code: def precompute_edge_matrices(adjacency, hparams):
    
    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
docs: 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.