ggnn version of the MPNN from Gilmer et al.

  Let B be the number of batches.
  Let D be the size of the node hidden states.
  Let K be the size of the attention keys/queries.
  Let V be the size of the output of the ggnn.
  Let T be the number of transforms / edge types.

  Args:
    node_states: The value Tensor of shape [B, T, N, D].
    adjacency_matrix: A Tensor of shape [B, N, N, T]. An entry at
      indices b, i, j, k is the indicator of the edge from node j to node i in
      batch b. A standard adjacency matrix will only have values of one, while a
      mutigraph may have larger integer values.
    num_edge_types: An integer specifying number of edge types.
    total_value_depth: An integer (V)
    name: A string.

  Returns:
    A Tensor of shape [B, N, V] storing the result of computing attention
    weights using the queries and keys and combining the values according to
    those weights.

  Raises:
    ValueError: if num_transforms doesn't equal num_edge_types and not using
      weighted sum.

def ggnn_fast_dense(node_states,
                    adjacency_matrix,
                    num_edge_types,
                    total_value_depth,
                    name=None):
  """ggnn version of the MPNN from Gilmer et al.

  Let B be the number of batches.
  Let D be the size of the node hidden states.
  Let K be the size of the attention keys/queries.
  Let V be the size of the output of the ggnn.
  Let T be the number of transforms / edge types.

  Args:
    node_states: The value Tensor of shape [B, T, N, D].
    adjacency_matrix: A Tensor of shape [B, N, N, T]. An entry at
      indices b, i, j, k is the indicator of the edge from node j to node i in
      batch b. A standard adjacency matrix will only have values of one, while a
      mutigraph may have larger integer values.
    num_edge_types: An integer specifying number of edge types.
    total_value_depth: An integer (V)
    name: A string.

  Returns:
    A Tensor of shape [B, N, V] storing the result of computing attention
    weights using the queries and keys and combining the values according to
    those weights.

  Raises:
    ValueError: if num_transforms doesn't equal num_edge_types and not using
      weighted sum.
  """
  # between the same nodes (with only one edge of each type. adjacency_matrix
  # will need to be converted to shape [B, T, N, N].
  with tf.variable_scope(
      name,
      default_name="ggnn_fast_dense",
      values=[node_states, adjacency_matrix, num_edge_types]):
    nodes_shape = common_layers.shape_list(node_states)
    v = _compute_edge_transforms(node_states,
                                 total_value_depth,
                                 num_edge_types,
                                 name="v_mpnn")
    v = tf.reshape(v, [nodes_shape[0], nodes_shape[1], num_edge_types,
                       total_value_depth
                      ])  # Shape [B, N, T, V].
    v = tf.transpose(v, [0, 2, 1, 3])  # Shape [B, T, N, V].

    # Rearranging the dimensions to match the shape of all_edge_logits.
    edge_vectors = tf.transpose(adjacency_matrix, [0, 3, 1, 2])
    output = compute_values(edge_vectors, v)
    return output