<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 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])