class_docstr: A callback to rewrite repeated scatters with a batched gather.
class_docstr_tok: ["A","callback","to","rewrite","repeated","scatters","with","a","batched","gather","."]
docstr: Replace the detected scatter loop with the following PyTorch code
  
  indices_per_level = []
  for level in range(num_scales):
      idx_in_level = torch.where(levels == level)[0]
      indices_per_leve.append(idx_in_level)
  
  stacked_features = torch.cat(roi_align_results, dim=0)
  stacked_indices = torch.cat(indices_per_level, dim=0)
  argsort_indices = torch.argort(stacked_indices)
  return stacked_features[argsort_indices, :]
docstr_tok: ["Replace","the","detected","scatter","loop","with","the","following","PyTorch","code","indices_per_level","=","[","]","for","level","in","range","(","num_scales",")",":","idx_in_level","=","torch.where","(","levels","==","level",")","[","0","]","indices_per_leve.append","(","idx_in_level",")","stacked_features","=","torch.cat","(","roi_align_results",",","dim=0",")","stacked_indices","=","torch.cat","(","indices_per_level",",","dim=0",")","argsort_indices","=","torch.argort","(","stacked_indices",")","return","stacked_features","[","argsort_indices",",",":","]"]
code: def convert_scatter_to_gather(self, levels, roi_align_results):
      """Replace the detected scatter loop with the following PyTorch code
  
      indices_per_level = []
      for level in range(num_scales):
          idx_in_level = torch.where(levels == level)[0]
          indices_per_leve.append(idx_in_level)
  
      stacked_features = torch.cat(roi_align_results, dim=0)
      stacked_indices = torch.cat(indices_per_level, dim=0)
      argsort_indices = torch.argort(stacked_indices)
      return stacked_features[argsort_indices, :]
      """
  
      # Collect inidices and concat them
      indices_per_level = []
      for i in range(self.num_scales):
          equal = op.equal(levels, expr.const(i, dtype="int64"))
          argwhere = op.argwhere(equal)
          split = op.split(argwhere, indices_or_sections=1, axis=1)
          squeeze = op.squeeze(split[0], axis=[1])
          indices = op.cast(squeeze, dtype="int64")
          indices_per_level.append(indices)
  
      indices_concat = op.concatenate(indices_per_level, 0)
  
      # Concat roi align results per level, and argsort indices
      # To prepare for a batched gather
      roi_align_results_concat = op.concatenate(roi_align_results, 0)
      argsort_indices = op.cast(
          op.argsort(indices_concat), dtype="int64"
      )
  
      # Permute rows by argsorted indices
      permuted = op.take(
          roi_align_results_concat, argsort_indices, axis=0
      )
  
      return op.reshape(permuted, [0, -1, 1, 1])
  
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 ScatterRewrite obj's underlying member method code with example usage: obj.convert_scatter_to_gather(levels, roi_align_results) and returns: op