code: 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])
  
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 6 for calling a global function with example usage: sparse_message_pass_batched(node_states, adjacency_matrices, num_edge_types, hidden_size, use_bias, average_aggregation, name) and returns: tf