<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_batched(node_states, adjacency_matrices, num_edge_types, hidden_size, use_bias=True, average_aggregation=False, name="sparse_ggnn_batched"):
 """Identical to sparse_ggnn except that each input has a batch dimension. B = The batch size. 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: [B, N, H] adjacency_matrices: Adjacency matrices of directed edges for each edge type and batch. Shape: [B, N, N, T] (sparse). num_edge_types: The number of edge types. T. hidden_size: The size of the hidden layer. 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 round of message-passing of shape [B, N, H]. """


  b, n = tf.shape(node_states)[0], tf.shape(node_states)[1]

  # Flatten the batch dimension of the node states.
  node_states = tf.reshape(node_states, [b*n, hidden_size])

  # Flatten the batch dimension of the adjacency matrices.
  indices = adjacency_matrices.indices
  new_index2 = indices[:, 3]  # The edge type dimension.

  # Offset N x N adjacency matrix by the batch number in which it appears.
  new_index0 = indices[:, 1] + indices[:, 0] * tf.cast(n, tf.int64)
  new_index1 = indices[:, 2] + indices[:, 0] * tf.cast(n, tf.int64)

  # Combine these indices as triples.
  new_indices = tf.stack([new_index0, new_index1, new_index2], axis=1)

  # Build the new sparse matrix.
  new_shape = [tf.cast(b*n, tf.int64), tf.cast(b*n, tf.int64), num_edge_types]
  adjacency_matrices = tf.SparseTensor(indices=new_indices,
                                       values=adjacency_matrices.values,
                                       dense_shape=new_shape)

  # Run a message-passing step and return the result with the batch dimension.
  node_states = sparse_message_pass(
      node_states,
      adjacency_matrices,
      num_edge_types,
      hidden_size,
      use_bias=use_bias,
      average_aggregation=average_aggregation,
      name=name)
  return tf.reshape(node_states, [b, n, hidden_size])