code: def to_dense_adj(edge_index, batch, edge_attr=None):
      assert batch is not None
      batch_size = batch[-1] + 1
  
      num_nodes = tf.scatter_nd(
          tf.reshape(batch, [-1, 1]),
          tf.ones(batch.shape, dtype=tf.int32),
          shape=[batch_size],
      )
  
      cum_nodes = tf.concat(
          [tf.zeros(1, dtype=tf.int32), tf.cumsum(num_nodes, axis=0)],
          axis=0,
      )
      max_num_nodes = tf.reduce_max(num_nodes, axis=0)
  
      shape = [batch_size, max_num_nodes, max_num_nodes]
      if edge_attr is not None:
          edge_attr_shape = edge_attr.shape.as_list()
          shape = shape + edge_attr_shape[1:]
  
      edge_index_0 = tf.gather(
          batch, edge_index[0]
      )  # [0, 0, 0, 1, 1, 1]
  
      cum_index = tf.gather(
          cum_nodes, edge_index_0
      )  # [0, 0, 0, 2, 2, 2]
  
      edge_index_1 = edge_index[0] - cum_index  # [0, 0, 1, 0, 1, 2]
      edge_index_2 = edge_index[1] - cum_index  # [0, 1, 0, 1, 2, 0]
  
      edge_index_0 = tf.reshape(edge_index_0, [-1, 1])
      edge_index_1 = tf.reshape(edge_index_1, [-1, 1])
      edge_index_2 = tf.reshape(edge_index_2, [-1, 1])
  
      idx = tf.concat(
          [edge_index_0, edge_index_1, edge_index_2], axis=-1
      )
  
      if edge_attr is None:
          edge_attr = tf.ones(edge_index.shape[-1])
      adj = tf.scatter_nd(idx, edge_attr, shape)
  
      return adj
  
label_desc: Use image node_id 1 for calling a global function with example usage: to_dense_adj(edge_index, batch, edge_attr) and returns: adj