code: def compute_mpnn_qkv(
      node_states, total_key_depth, total_value_depth, num_transforms
  ):
      """Computes query, key and value for edge matrices.
  
      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 queries, keys, and values for attention.
      * For each node N_i in the graph, a query Q_i of size K is computed. This
        query is used to determine the relative weights to give to each of the
        node's incoming edges.
      * 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 with shape [B, N, D].
        total_key_depth: an integer (K).
        total_value_depth: an integer (V).
        num_transforms: a integer specifying number of transforms (T). This is
          typically the number of edge types.
      Returns:
        q: The attention queries for each destination node (shape [B, N, K]).
        k: The attention keys for each node and edge type (shape [B, N*T, K]).
        v: The attention values for each node and edge type (shape [B, N*T, V]).
      """
  
      # node_states is initially a tensor with shape [B, N, D]. The call to dense
      # creates a D x K kernel that serves as a fully-connected layer.
      #
      # For each possible batch b and node n in the first two dimensions of
      # node_states, the corresponding size-D vector (the third dimension of
      # node_states) is the hidden state for node n in batch b. Each of these size-D
      # vectors is multiplied by the kernel to produce an attention query of size K.
      # The result is a tensor of size [B, N, K] containing the attention queries
      # for each node in each batch.
      q = common_layers.dense(
          node_states, total_key_depth, use_bias=False, name="q_mpnn"
      )
  
      # Creates the attention keys in a manner similar to the process of creating
      # the attention queries. One key is created for each type of outgoing edge the
      # corresponding node might have, meaning k will have shape [B, N, K*T].
      k = _compute_edge_transforms(
          node_states, total_key_depth, num_transforms, name="k_mpnn"
      )
      v = _compute_edge_transforms(
          node_states, total_value_depth, num_transforms, name="v_mpnn"
      )
  
      return q, k, v
  
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 5 for calling a global function with example usage: compute_mpnn_qkv(node_states, total_key_depth, total_value_depth, num_transforms) and returns: q,  k,  v