code: def topk_(input, k, axis=1, largest=True):
      x = -input if largest else input
      if axis == 0:
          row_index = np.arange(input.shape[1 - axis])
          if (
              k == x.shape[0]
          ):  # argpartition requires index < len(input)
              topk_index = np.argpartition(x, k - 1, axis=axis)[0:k, :]
          else:
              topk_index = np.argpartition(x, k, axis=axis)[0:k, :]
  
          topk_data = x[topk_index, row_index]
  
          topk_index_sort = np.argsort(topk_data, axis=axis)
          topk_data_sort = topk_data[topk_index_sort, row_index]
          topk_index_sort = topk_index[0:k, :][
              topk_index_sort, row_index
          ]
      else:
          column_index = np.arange(x.shape[1 - axis])[:, None]
          topk_index = np.argpartition(x, k, axis=axis)[:, 0:k]
          topk_data = x[column_index, topk_index]
          topk_data = -topk_data if largest else topk_data
          topk_index_sort = np.argsort(topk_data, axis=axis)
          topk_data_sort = topk_data[column_index, topk_index_sort]
          topk_index_sort = topk_index[:, 0:k][
              column_index, topk_index_sort
          ]
  
      return topk_data_sort, topk_index_sort
  
file_path: PaddleDetection/ppdet/data/transform/atss_assigner.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  from ppdet.utils.logger import setup_logger
label_desc: Use image node_id 2 for calling a global function with example usage: topk_(input, k, axis, largest) and returns: topk_data_sort,  topk_index_sort