code: def _to_nd_indices(indices):
      """Returns indices used for tf.gather_nd or tf.scatter_nd.
  
      Args:
        indices: A `Tensor` of shape [batch_size, size] with integer values. The
          values are the indices of another `Tensor`. For example, `indices` is the
          output of tf.argsort or tf.math.top_k.
  
      Returns:
        A `Tensor` with shape [batch_size, size, 2] that can be used by tf.gather_nd
        or tf.scatter_nd.
  
      """
      indices.get_shape().assert_has_rank(2)
      batch_ids = tf.ones_like(indices) * tf.expand_dims(
          tf.range(tf.shape(input=indices)[0]), 1
      )
      return tf.stack([batch_ids, indices], axis=-1)
  
file_path: ranking/tensorflow_ranking/python/utils.py
imports: from typing import Callable, Dict, Tuple
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: _to_nd_indices(indices) and returns: tf