code: def scatter_roi_align_result_pattern(
      levels, roi_align_results, num_scales
  ):
      """Detect the Relay subgraph corresponding to the following PyTorch code
  
      first_result = roi_align_results[0]
      dtype, device = first_result.dtype, first_result.device
      res = torch.zeros((levels.size(0), first_result.size(1),
                         first_result.size(2), first_result.size(3)),
                        dtype=dtype, device=device)
      for level in range(len(roi_align_results)):
          index = torch.where(levels == level)[0].view(-1, 1, 1, 1)
          index = index.expand(index.size(0),
                               roi_align_results[level].size(1),
                               roi_align_results[level].size(2),
                               roi_align_results[level].size(3))
          res = res.scatter(0, index, roi_align_results[level])
      return res
      """
  
      def do_where(levels, _):
          idx_in_level = is_op("argwhere")(
              is_op("equal")(levels, is_constant())
          )
          idx_in_level = is_op("split")(idx_in_level)
          idx_in_level = is_tuple_get_item(idx_in_level, 0)
          idx_in_level = is_op("squeeze")(idx_in_level)
          idx_in_level = is_tuple_get_item(is_tuple([idx_in_level]), 0)
          return idx_in_level
  
      scatter_res = wildcard()
  
      for i in range(num_scales):
          # index = torch.where(levels == level)[0].view(-1, 1, 1, 1)
          scatter_indices = do_where(levels, i)
          scatter_indices = is_op("reshape")(scatter_indices)
  
          # index = index.expand(index.size(0),
          #                      unmerged_results[level].size(1),
          #                      unmerged_results[level].size(2),
          #                      unmerged_results[level].size(3))
          scatter_indices = is_op("repeat")(scatter_indices)
          scatter_indices = is_op("repeat")(scatter_indices)
          scatter_indices = is_op("repeat")(scatter_indices)
  
          scatter_res = is_op("scatter_elements")(
              scatter_res, scatter_indices, roi_align_results[i]
          )
  
      return is_op("reshape")(scatter_res)
  
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 6 for calling a global function with example usage: scatter_roi_align_result_pattern(levels, roi_align_results, num_scales) and returns: is_op, idx_in_level