code: def make_edge_vectors(
      adjacency_matrix, num_edge_types, depth, name=None
  ):
      """Gets edge vectors for the edge types in the adjacency matrix.
  
      Args:
        adjacency_matrix: A [batch, num_nodes, num_nodes, num_edge_types] tensor.
        num_edge_types: Number of different edge types
        depth: Number of channels
        name: A optional string name for scoping
      Returns:
        A [batch, num_nodes, num_nodes, depth] vector of tensors
      """
      with tf.variable_scope(name, default_name="edge_vectors"):
          att_adj_vectors_shape = [num_edge_types, depth]
          adjacency_matrix_shape = common_layers.shape_list(
              adjacency_matrix
          )
          adj_vectors = tf.get_variable(
              "adj_vectors",
              att_adj_vectors_shape,
              initializer=tf.random_normal_initializer(
                  0, depth**-0.5
              ),
          ) * (depth**0.5)
  
          att_adj_vectors = tf.matmul(
              tf.reshape(
                  tf.to_float(adjacency_matrix), [-1, num_edge_types]
              ),
              adj_vectors,
          )
          # Reshape to be [batch, num_nodes, num_nodes, depth].
          att_adj_vectors = tf.reshape(
              att_adj_vectors,
              [
                  adjacency_matrix_shape[0],
                  adjacency_matrix_shape[1],
                  adjacency_matrix_shape[2],
                  depth,
              ],
          )
          return att_adj_vectors
  
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 2 for calling a global function with example usage: make_edge_vectors(adjacency_matrix, num_edge_types, depth, name) and returns: att_adj_vectors