code: def sparse_message_pass(
      node_states,
      adjacency_matrices,
      num_edge_types,
      hidden_size,
      use_bias=True,
      average_aggregation=False,
      name="sparse_ggnn",
  ):
      """One message-passing step for a GNN with a sparse adjacency matrix.
  
      Implements equation 2 (the message passing step) in
      [Li et al. 2015](https://arxiv.org/abs/1511.05493).
  
      N = The number of nodes in each batch.
      H = The size of the hidden states.
      T = The number of edge types.
  
      Args:
        node_states: Initial states of each node in the graph. Shape is [N, H].
        adjacency_matrices: Adjacency matrix of directed edges for each edge
          type. Shape is [N, N, T] (sparse tensor).
        num_edge_types: The number of edge types. T.
        hidden_size: The size of the hidden state. H.
        use_bias: Whether to use bias in the hidden layer.
        average_aggregation: How to aggregate the incoming node messages. If
          average_aggregation is true, the messages are averaged. If it is false,
          they are summed.
        name: (optional) The scope within which tf variables should be created.
  
      Returns:
        The result of one step of Gated Graph Neural Network (GGNN) message passing.
        Shape: [N, H]
      """
      n = tf.shape(node_states)[0]
      t = num_edge_types
      incoming_edges_per_type = tf.sparse_reduce_sum(
          adjacency_matrices, axis=1
      )
  
      # Convert the adjacency matrix into shape [T, N, N] - one [N, N] adjacency
      # matrix for each edge type. Since sparse tensor multiplication only supports
      # two-dimensional tensors, we actually convert the adjacency matrix into a
      # [T * N, N] tensor.
      adjacency_matrices = tf.sparse_transpose(
          adjacency_matrices, [2, 0, 1]
      )
      adjacency_matrices = tf.sparse_reshape(
          adjacency_matrices, [t * n, n]
      )
  
      # Multiply the adjacency matrix by the node states, producing a [T * N, H]
      # tensor. For each (edge type, node) pair, this tensor stores the sum of
      # the hidden states of the node's neighbors over incoming edges of that type.
      messages = tf.sparse_tensor_dense_matmul(
          adjacency_matrices, node_states
      )
  
      # Rearrange this tensor to have shape [N, T * H]. The incoming states of each
      # nodes neighbors are summed by edge type and then concatenated together into
      # a single T * H vector.
      messages = tf.reshape(messages, [t, n, hidden_size])
      messages = tf.transpose(messages, [1, 0, 2])
      messages = tf.reshape(messages, [n, t * hidden_size])
  
      # Run each of those T * H vectors through a linear layer that produces
      # a vector of size H. This process is equivalent to running each H-sized
      # vector through a separate linear layer for each edge type and then adding
      # the results together.
      #
      # Note that, earlier on, we added together all of the states of neighbors
      # that were connected by edges of the same edge type. Since addition and
      # multiplying by a linear layer are commutative, this process was equivalent
      # to running each incoming edge through a linear layer separately and then
      # adding everything at the end.
      with tf.variable_scope(name, default_name="sparse_ggnn"):
          final_node_states = common_layers.dense(
              messages, hidden_size, use_bias=False
          )
  
          # Multiply the bias by for each edge type by the number of incoming nodes
          # of that edge type.
          if use_bias:
              bias = tf.get_variable(
                  "bias", initializer=tf.zeros([t, hidden_size])
              )
              final_node_states += tf.matmul(
                  incoming_edges_per_type, bias
              )
  
          if average_aggregation:
              incoming_edges = tf.reduce_sum(
                  incoming_edges_per_type, -1, keepdims=True
              )
              incoming_edges = tf.tile(incoming_edges, [1, hidden_size])
              final_node_states /= incoming_edges + 1e-7
  
      return tf.reshape(final_node_states, [n, hidden_size])
  
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 7 for calling a global function with example usage: sparse_message_pass(node_states, adjacency_matrices, num_edge_types, hidden_size, use_bias, average_aggregation, name) and returns: tf