code: def test_partition_indices_execution(setup):
      # only 1 chunk when axis = -1
      raw = np.random.rand(100, 10)
      x = tensor(raw, chunk_size=10)
  
      kth = [2, 5, 9]
      r = partition(x, kth, return_index=True)
  
      pr, pi = r.execute().fetch()
      np.testing.assert_array_equal(
          pr, np.take_along_axis(raw, pi, axis=-1)
      )
      np.testing.assert_array_equal(np.sort(raw)[:, kth], pr[:, kth])
  
      x = tensor(raw, chunk_size=(22, 4))
  
      r = partition(x, kth, return_index=True)
  
      pr, pi = r.execute().fetch()
      np.testing.assert_array_equal(
          pr, np.take_along_axis(raw, pi, axis=-1)
      )
      np.testing.assert_array_equal(np.sort(raw)[:, kth], pr[:, kth])
  
      raw = np.random.rand(100)
  
      x = tensor(raw, chunk_size=23)
  
      r = partition(x, kth, axis=0, return_index=True)
  
      pr, pi = r.execute().fetch()
      np.testing.assert_array_equal(
          pr, np.take_along_axis(raw, pi, axis=-1)
      )
      np.testing.assert_array_equal(np.sort(raw)[kth], pr[kth])
  
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 36 for calling a global function with example usage: test_partition_indices_execution(setup) without return types