code: def batched_nms_pattern(
      boxes, scores, idxs, iou_threshold, num_boxes, indices
  ):
      """A pattern to detect batched_nms function in torchvision
  
      The inputs to this function, boxes, scores, idxs, iou_threshold are wildcard
      patterns which can be used later in the rewriting to extract matched Relay fragments.
  
      We want to detect the following PyTorch code snippet:
  
      def batched_nms(boxes, scores, idxs, iou_threshold):
          max_coordinate = boxes.max()
          offsets = idxs.to(boxes) * (max_coordinate + torch.tensor(1).to(boxes))
          boxes_for_nms = boxes + offsets[:, None]
          keep = nms(boxes_for_nms, scores, iou_threshold)
          return keep
  
      Here is how PyTorch frontend lowers above PyTorch code. For simplicity, Relay ops for
      dealing with dynamic strided_slice are omitted. %num_boxes, %indices are complex
      expressions, but since we can use the wildcard part for them, we do not need to construct
      their patterns.
  
      %2 = expand_dims(%scores, axis=-1);
      %3 = cast(%idxs, dtype="float32");
      %4 = max(%boxes);
      %5 = add(%4, 1f);
      %6 = multiply(%3, %5);
      %7 = strided_slice(%6, begin=[0], end=[4507], strides=[1]);
      %8 = expand_dims(%7, axis=1);
      %9 = add(%boxes, %8);
      %10 = (%2, %9);
      %11 = concatenate(%10, axis=-1);
      %12 = expand_dims(%11, axis=0);
      ...
      ...
      %17 = vision.non_max_suppression(%12, %num_boxes, %indices, -1, 0.7f, ...);
  
      """
      one = is_constant()
  
      # Equivalent PyTorch code from above snippet
      # offsets = idxs.to(boxes) * (max_coordinate + torch.tensor(1).to(boxes))
      cast = is_op("cast")(idxs)
      mx = is_op("max")(boxes)
      add = is_op("add")(mx, one)
      mul = is_op("multiply")(cast, add)
  
      shape_of = is_op("shape_of")(mul)
      cast = is_op("cast")(shape_of)
  
      # Add offsets to the boxes
      expand_dims = is_op("expand_dims")(mul)
      add = is_op("add")(boxes, expand_dims)
  
      # The rest of patterns correspond to the PyTorch frontend conversion
      # function for torchvision::nms
      score_expand_dims = is_op("expand_dims")(scores)
      tup = is_tuple([score_expand_dims, add])
      concat = is_op("concatenate")(tup)
      data = is_op("expand_dims")(concat)
  
      return is_op("vision.non_max_suppression")(
          data, num_boxes, indices, is_constant(), iou_threshold
      )
  
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 4 for calling a global function with example usage: batched_nms_pattern(boxes, scores, idxs, iou_threshold, num_boxes, indices) and returns: is_op