code: def _get_q_slice(q, k, ind, b=None, batch_shape=None):
      """Returns `q1[i]` or `q0[j]` for a batch of indices `i` or `j`."""
      q_ind = tf.concat(
          [ind, tf.expand_dims(tf.gather_nd(k, ind), -1)], axis=1
      )
      b_updates = tf.gather_nd(q, q_ind)
      if b is None:
          return tf.scatter_nd(ind, b_updates, batch_shape)
      return tf.tensor_scatter_nd_update(b, ind, b_updates)
  
file_path: probability/tensorflow_probability/python/stats/ranking.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
label_desc: Use image node_id 2 for calling a global function with example usage: _get_q_slice(q, k, ind, b, batch_shape) and returns: tf, tf