class_docstr: A callback to rewrite nms to exploit max_out_size parameter.
class_docstr_tok: ["A","callback","to","rewrite","nms","to","exploit","max_out_size","parameter","."]
docstr: Use the detected post NMS topk parameter in NMS op.
docstr_tok: ["Use","the","detected","post","NMS","topk","parameter","in","NMS","op","."]
code: def rewrite_batch_nms_with_max_out_size(
      self,
      cond,
      true_branch,
      data,
      valid_count,
      indices,
      iou_threshold,
      post_nms_topk,
  ):
      """Use the detected post NMS topk parameter in NMS op."""
      nms_ret = op.vision.non_max_suppression(
          data=data,
          valid_count=valid_count,
          indices=indices,
          max_output_size=post_nms_topk,
          iou_threshold=iou_threshold,
          force_suppress=False,
          top_k=-1,
          coord_start=2,
          score_index=1,
          id_index=0,
          return_indices=True,
          invalid_to_bottom=False,
      )
  
      size = op.squeeze(nms_ret[1], axis=[1])
      data_slice = op.squeeze(nms_ret[0], axis=[0])
  
      ret = op.strided_slice(
          data_slice, begin=expr.const([0]), end=size, slice_mode="size"
      )
  
      nms_result = op.cast(ret, "int64")
  
      return expr.If(cond, true_branch, nms_result)
  
imports: from ..None import expr
  from ..None import op
  from ..dataflow_pattern import wildcard, is_constant, is_op, rewrite, is_tuple, is_tuple_get_item, is_if, DFPatternCallback
label_desc: Use image node_id 2 for calling the PostNMSTopKRewrite obj's underlying member method code with example usage: obj.rewrite_batch_nms_with_max_out_size(cond, true_branch, data, valid_count, indices, iou_threshold, post_nms_topk) and returns: expr