code: def get_valid_counts_ir(
      data, valid_indices, valid_boxes, out, out_indices
  ):
      """Low level IR to get valid count of bounding boxes
      given a score threshold. Also prepares to move valid boxes to the
      top of input data.
  
      Parameters
      ----------
      data : Buffer
          Input data. 3-D Buffer with shape [batch_size, num_anchors, elem_length].
  
      valid_indices: Buffer
          2D Buffer of flag indicating valid data with shape [batch_size, num_anchors].
  
      Returns
      -------
      out : Buffer
          Sorted valid boxes
  
      out_indices : Buffer
          Incidices of valid boxes in original data
      """
      batch_size = data.shape[0]
      num_anchors = data.shape[1]
      elem_length = data.shape[2]
  
      ib = tvm.tir.ir_builder.create()
  
      data = ib.buffer_ptr(data)
      valid_indices = ib.buffer_ptr(valid_indices)
      valid_boxes = ib.buffer_ptr(valid_boxes)
  
      out = ib.buffer_ptr(out)
      out_indices = ib.buffer_ptr(out_indices)
      one = tvm.tir.const(1, dtype=out.dtype)
  
      max_threads = int(
          tvm.target.Target.current(allow_none=False).max_num_threads
      )
      nthread_tx = max_threads
      nthread_bx = num_anchors // max_threads + 1
      nthread_by = batch_size
      with ib.new_scope():
          tx = te.thread_axis("threadIdx.x")
          bx = te.thread_axis("blockIdx.x")
          by = te.thread_axis("blockIdx.y")
          ib.scope_attr(tx, "thread_extent", nthread_tx)
          ib.scope_attr(bx, "thread_extent", nthread_bx)
          ib.scope_attr(by, "thread_extent", nthread_by)
          tid = bx * max_threads + tx
          with ib.if_scope(tid < num_anchors):
              i = by
              j = tid
              with ib.for_range(0, elem_length) as k:
                  out[(i * num_anchors + j) * elem_length + k] = -one
              out_indices[i * num_anchors + j] = -1
      with ib.new_scope():
          tx = te.thread_axis("threadIdx.x")
          bx = te.thread_axis("blockIdx.x")
          by = te.thread_axis("blockIdx.y")
          ib.scope_attr(tx, "thread_extent", nthread_tx)
          ib.scope_attr(bx, "thread_extent", nthread_bx)
          ib.scope_attr(by, "thread_extent", nthread_by)
          tid = bx * max_threads + tx
          with ib.if_scope(tid < num_anchors):
              i = by
              j = tid
              with ib.if_scope(valid_boxes[i, tid] > 0):
                  with ib.for_range(0, elem_length) as k:
                      out[
                          (i * num_anchors + valid_indices[i, tid])
                          * elem_length
                          + k
                      ] = data[(i * num_anchors + j) * elem_length + k]
                  out_indices[
                      i * num_anchors + valid_indices[i, tid]
                  ] = j
      return ib.get()
  
imports: import tvm
  from tvm import te
  from tvm.contrib import nvcc
  from tvm.contrib.thrust import can_use_thrust, can_use_rocthrust
  from tvm.ir import register_intrin_lowering
  from tvm.tir import if_then_else
  from .sort import argsort, argsort_thrust
  from .scan import exclusive_scan
  from ..utils import ceil_div
  from ..math import cast
  from ..transform import reshape
  from ..vision.nms_util import calculate_overlap, binary_search, collect_selected_indices, collect_selected_indices_and_scores, run_all_class_nms
label_desc: Use image node_id 5 for calling a global function with example usage: get_valid_counts_ir(data, valid_indices, valid_boxes, out, out_indices) and returns: ib