code: def _form_group_indices_nd(
      is_valid, group_size, shuffle=False, seed=None
  ):
      """Forms the indices for groups for gather_nd or scatter_nd.
  
      Args:
        is_valid: A boolean `Tensor` for entry validity with shape [batch_size,
          list_size].
        group_size: An scalar int `Tensor` for the number of examples in a group.
        shuffle: A boolean that indicates whether valid indices should be shuffled
          when forming group indices.
        seed: Random seed for shuffle.
  
      Returns:
        A tuple of Tensors (indices, mask). The first has shape [batch_size,
        num_groups, group_size, 2] and it can be used in gather_nd or scatter_nd for
        group features. The second has the shape of [batch_size, num_groups] with
        value True for valid groups.
      """
      with tf.compat.v1.name_scope(name="form_group_indices"):
          is_valid = tf.convert_to_tensor(value=is_valid)
          batch_size, list_size = tf.unstack(tf.shape(input=is_valid))
          num_valid_entries = tf.reduce_sum(
              input_tensor=tf.cast(is_valid, dtype=tf.int32), axis=1
          )
          rw_indices, mask = _rolling_window_indices(
              list_size, group_size, num_valid_entries
          )
          # Valid indices of the tensor are shuffled and put on the top.
          # [batch_size, list_size, 2]. A deterministic op-level seed is set mainly
          # for unittest purpose. We can find a better way to avoid setting this seed
          # explicitly.
          shuffled_indices = utils.organize_valid_indices(
              is_valid, shuffle=shuffle, seed=seed
          )
          # Construct indices for gather_nd.
          # [batch_size, num_groups, group_size, 2]
          group_indices_nd = tf.expand_dims(rw_indices, axis=3)
          group_indices_nd = tf.concat(
              [
                  tf.reshape(tf.range(batch_size), [-1, 1, 1, 1])
                  * tf.ones_like(group_indices_nd),
                  group_indices_nd,
              ],
              3,
          )
  
          indices = tf.gather_nd(shuffled_indices, group_indices_nd)
          return indices, mask
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import abc
  import six
  import tensorflow
  from tensorflow import estimator
  from tensorflow.python.util import function_utils
  from tensorflow_ranking.python import feature
  from tensorflow_ranking.python import utils
label_desc: Use image node_id 3 for calling a global function with example usage: _form_group_indices_nd(is_valid, group_size, shuffle, seed) and returns: indices,  mask