code: def segment_top_k(x, i, ratio):
      """
      Returns indices to get the top K values in x segment-wise, according to
      the segments defined in I. K is not fixed, but it is defined as a ratio of
      the number of elements in each segment.
      :param x: a rank 1 Tensor;
      :param i: a rank 1 Tensor with segment IDs for x;
      :param ratio: float, ratio of elements to keep for each segment;
      :return: a rank 1 Tensor containing the indices to get the top K values of
      each segment in x.
      """
      i = tf.cast(i, tf.int32)
      n = tf.shape(i)[0]
      n_nodes = tf.math.segment_sum(tf.ones_like(i), i)
      batch_size = tf.shape(n_nodes)[0]
      n_nodes_max = tf.reduce_max(n_nodes)
      cumulative_n_nodes = tf.concat(
          (tf.zeros(1, dtype=n_nodes.dtype), tf.cumsum(n_nodes)[:-1]), 0
      )
      index = tf.range(n)
      index = (index - tf.gather(cumulative_n_nodes, i)) + (
          i * n_nodes_max
      )
  
      dense_x = tf.zeros(batch_size * n_nodes_max, dtype=x.dtype) - 1e20
      dense_x = tf.tensor_scatter_nd_update(dense_x, index[:, None], x)
      dense_x = tf.reshape(dense_x, (batch_size, n_nodes_max))
  
      perm = tf.argsort(dense_x, direction="DESCENDING")
      perm = perm + cumulative_n_nodes[:, None]
      perm = tf.reshape(perm, (-1,))
  
      k = tf.cast(
          tf.math.ceil(ratio * tf.cast(n_nodes, tf.float32)), i.dtype
      )
  
      # This costs more memory
      # to_rep = tf.tile(tf.constant([1., 0.]), (batch_size,))
      # rep_times = tf.reshape(tf.concat((k[:, None], (n_nodes_max - k)[:, None]), -1), (-1,))
      # mask = ops.repeat(to_rep, rep_times)
      # perm = tf.boolean_mask(perm, mask)
  
      # This is slower
      r_range = tf.ragged.range(k).flat_values
      r_delta = ops.repeat(tf.range(batch_size) * n_nodes_max, k)
      mask = r_range + r_delta
      perm = tf.gather(perm, mask)
  
      return perm
  
file_path: spektral/spektral/layers/pooling/topk_pool.py
imports: import tensorflow
  from tensorflow.keras import backend
  from spektral.layers import ops
  from spektral.layers.pooling.src import SRCPool
label_desc: Use image node_id 1 for calling a global function with example usage: segment_top_k(x, i, ratio) and returns: perm