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.
      """
      rt = tf.RaggedTensor.from_value_rowids(x, I)
      row_lengths = rt.row_lengths()
      dense = rt.to_tensor(default_value=-np.inf)
      indices = tf.cast(
          tf.argsort(dense, direction="DESCENDING"), tf.int64
      )
      row_starts = tf.cast(rt.row_starts(), tf.int64)
      indices = indices + tf.expand_dims(row_starts, 1)
      row_lengths = tf.cast(
          tf.math.ceil(ratio * tf.cast(row_lengths, tf.float32)),
          tf.int32,
      )
      return tf.RaggedTensor.from_tensor(indices, row_lengths).values
  
imports: import numpy
  import tensorflow
  from tensorflow.keras import backend
label_desc: Use image node_id 4 for calling a global function with example usage: segment_top_k(x, I, ratio) and returns: tf