code: def collect_selected_indices_and_scores(
      selected_indices,
      selected_scores,
      num_detections,
      row_offsets,
      num_total_detections,
      ir,
  ):
      """Collect selected indices and scores from the core NMS loop into one linear output
  
      Parameters
      ----------
      num_class : int
  
      selected_indices: tvm.te.Tensor
          2-D tensor with shape (batch_size * num_classes, num_boxes), representing the indices
          of selected boxes by the core NMS loop.
  
      selected_indices: tvm.te.Tensor
          2-D tensor with shape (batch_size * num_classes, num_boxes), representing the scores
          of selected boxes by the core NMS loop.
  
      num_detections tvm.te.Tensor
          2-D tensor with shape (batch_size, num_classes), representing
          the number of boxes selected by the core NMS loop, per batch and class
  
      row_offsets tvm.te.Tensor
          2-D tensor with shape (batch_size, num_classes), this should be the exclusive scan
          of num_detections along axis 1
  
      ir : function
          A function to generate IR for CPU or GPU, see its usage in vision/nms.py and cuda/nms.py
  
      Returns
      -------
      out : [tvm.te.Tensor, tvm.te.Tensor]
          The output is two tensors. The first is indices of size
          (batch_size, num_class* num_boxes, 2), and the second is scores of size
          (batch_size, num_class* num_boxes).
      """
      batch_size, num_class = row_offsets.shape
      num_boxes = selected_indices.shape[1]
      return te.extern(
          [
              (batch_size, num_class * num_boxes, 2),
              (batch_size, num_class * num_boxes),
          ],
          [
              selected_indices,
              selected_scores,
              num_detections,
              row_offsets,
              num_total_detections,
          ],
          lambda ins, outs: ir(
              ins[0], ins[1], ins[2], ins[3], ins[4], outs[0], outs[1]
          ),
          dtype=["int64", "float32"],
          name="collect_indices_and_scores",
          tag="collect_indices_and_scores",
      )
  
file_path: tvm/python/tvm/topi/vision/nms_util.py
imports: import tvm
  from tvm import te
label_desc: Use image node_id 5 for calling a global function with example usage: collect_selected_indices_and_scores(selected_indices, selected_scores, num_detections, row_offsets, num_total_detections, ir) and returns: te