code: 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
  
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 12 for calling a global function with example usage: precompute_edge_matrices(adjacency, hparams) and returns: edge_matrices