code: def test_argtopk(setup):
      # only 1 chunk when axis = -1
      raw = np.random.rand(100, 10)
      x = tensor(raw, chunk_size=20)
  
      pa = argtopk(x, 3, parallel_kind="tree")
  
      r = pa.execute().fetch()
      np.testing.assert_array_equal(
          np.sort(raw)[:, -1:-4:-1], np.take_along_axis(raw, r, axis=-1)
      )
  
      pa = argtopk(x, 3, parallel_kind="psrs")
  
      r = pa.execute().fetch()
      np.testing.assert_array_equal(
          np.sort(raw)[:, -1:-4:-1], np.take_along_axis(raw, r, axis=-1)
      )
  
      x = tensor(raw, chunk_size=(22, 4))
  
      pa = argtopk(x, 3, parallel_kind="tree")
  
      r = pa.execute().fetch()
      np.testing.assert_array_equal(
          np.sort(raw)[:, -1:-4:-1], np.take_along_axis(raw, r, axis=-1)
      )
  
      pa = argtopk(x, 3, parallel_kind="psrs")
  
      r = pa.execute().fetch()
      np.testing.assert_array_equal(
          np.sort(raw)[:, -1:-4:-1], np.take_along_axis(raw, r, axis=-1)
      )
  
      raw = np.random.rand(100)
  
      x = tensor(raw, chunk_size=23)
  
      pa = argtopk(x, 3, axis=0, parallel_kind="tree")
  
      r = pa.execute().fetch()
      np.testing.assert_array_equal(
          np.sort(raw, axis=0)[-1:-4:-1], raw[r]
      )
  
      pa = argtopk(x, 3, axis=0, parallel_kind="psrs")
  
      r = pa.execute().fetch()
      np.testing.assert_array_equal(
          np.sort(raw, axis=0)[-1:-4:-1], raw[r]
      )
  
imports: import numpy
  import scipy.sparse
  import pandas
  import pytest
  from ....None import dataframe
  from ....None import tensor
  from ....None import execute, fetch
  from ....tests.core import require_cupy
  from ...datasource import tensor, ones, zeros, arange
  from ..None import copyto, transpose, moveaxis, broadcast_to, broadcast_arrays, where, expand_dims, rollaxis, atleast_1d, atleast_2d, atleast_3d, argwhere, array_split, split, hsplit, vsplit, dsplit, roll, squeeze, diff, ediff1d, flip, flipud, fliplr, repeat, tile, isin, searchsorted, unique, sort, argsort, partition, argpartition, topk, argtopk, trapz, shape, to_gpu, to_cpu, swapaxes
label_desc: Use image node_id 41 for calling a global function with example usage: test_argtopk(setup) without return types