code: def _get_offsets(offset, topk_inds):
      b, _, _, n = offset.shape
      offset = tf.reshape(offset, [b, -1, n])
      offset = tf.gather(offset, topk_inds, batch_dims=1)
      return offset
  
imports: import gin
  import tensorflow
  from tensorflow_graphics.projects.points_to_3Dobjects.networks import networks
  import tensorflow_graphics.projects.points_to_3Dobjects.networks.custom_blocks
label_desc: Use image node_id 2 for calling a global function with example usage: _get_offsets(offset, topk_inds) and returns: offset