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
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]).
  
tensor2tensor/layers/message_passing_attention.py