code: def _compute_edge_transforms(
      node_states, depth, num_transforms, name="transform"
  ):
      """Helper function that computes transformation for keys and values.
  
      Let B be the number of batches.
      Let N be the number of nodes in the graph.
      Let D be the size of the node hidden states.
      Let K be the size of the attention keys/queries (total_key_depth).
      Let V be the size of the attention values (total_value_depth).
      Let T be the total number of transforms (num_transforms).
  
      Computes the transforms for keys or values for attention.
      * For each node N_j and edge type t, a key K_jt of size K is computed. When an
        edge of type t goes from node N_j to any other node, K_jt is the key that is
        in the attention process.
      * For each node N_j and edge type t, a value V_jt of size V is computed. When
        an edge of type t goes from node N_j to node N_i, Attention(Q_i, K_jt)
        produces a weight w_ijt. The message sent along this edge is w_ijt * V_jt.
  
      Args:
        node_states: A tensor of shape [B, L, D]
        depth: An integer (K or V)
        num_transforms: An integer (T),
        name: A name for the function
  
      Returns:
        x: A The attention keys or values for each node and edge type
          (shape [B, N*T, K or V])
      """
      node_shapes = common_layers.shape_list(node_states)
      x = common_layers.dense(
          node_states, depth * num_transforms, use_bias=False, name=name
      )
  
      batch = node_shapes[0]  # B.
      length = node_shapes[1]  # N.
  
      # Making the fourth dimension explicit by separating the vectors of size
      # K*T (in k) and V*T (in v) into two-dimensional matrices with shape [K, T]
      # (in k) and [V, T] in v.
      #
      x = tf.reshape(x, [batch, length, num_transforms, depth])
  
      # Flatten out the fourth dimension.
      x = tf.reshape(x, [batch, length * num_transforms, depth])
  
      return x
  
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 4 for calling a global function with example usage: _compute_edge_transforms(node_states, depth, num_transforms, name) and returns: x