code: def disjoint_signal_to_batch(X, I):
      """
      Converts a disjoint graph signal to batch node by zero-padding.
  
      :param X: Tensor, node features of shape (nodes, features).
      :param I: Tensor, graph IDs of shape `(n_nodes, )`;
      :return batch: Tensor, batched node features of shape (batch, N_max, n_node_features)
      """
      I = tf.cast(I, tf.int32)
      num_nodes = tf.math.segment_sum(tf.ones_like(I), I)
      start_index = tf.cumsum(num_nodes, exclusive=True)
      n_graphs = tf.shape(num_nodes)[0]
      max_n_nodes = tf.reduce_max(num_nodes)
      batch_n_nodes = tf.shape(I)[0]
      feature_dim = tf.shape(X)[-1]
  
      index = tf.range(batch_n_nodes)
      index = (index - tf.gather(start_index, I)) + (I * max_n_nodes)
      dense = tf.zeros(
          (n_graphs * max_n_nodes, feature_dim), dtype=X.dtype
      )
      dense = tf.tensor_scatter_nd_update(dense, index[..., None], X)
  
      batch = tf.reshape(dense, (n_graphs, max_n_nodes, feature_dim))
  
      return batch
  
imports: import tensorflow
  from tensorflow.keras import backend
label_desc: Use image node_id 1 for calling a global function with example usage: disjoint_signal_to_batch(X, I) and returns: batch