code: def test_partition_execution(setup):
      # only 1 chunk when axis = -1
      raw = np.random.rand(100, 10)
      x = tensor(raw, chunk_size=20)
  
      px = partition(x, [1, 8])
  
      res = px.execute().fetch()
      np.testing.assert_array_equal(res, np.partition(raw, [1, 8]))
  
      # 1-d chunk
      raw = np.random.rand(100)
      x = tensor(raw, chunk_size=20)
  
      kth = np.random.RandomState(0).randint(-100, 100, size=(10,))
      px = partition(x, kth)
  
      res = px.execute().fetch()
      np.testing.assert_array_equal(
          res[kth], np.partition(raw, kth)[kth]
      )
  
      # structured dtype
      raw = np.empty(100, dtype=[("id", np.int32), ("size", np.int64)])
      raw["id"] = np.random.randint(1000, size=100, dtype=np.int32)
      raw["size"] = np.random.randint(1000, size=100, dtype=np.int64)
      x = tensor(raw, chunk_size=20)
  
      px = partition(x, kth, order=["size", "id"])
  
      res = px.execute().fetch()
      np.testing.assert_array_equal(
          res[kth], np.partition(raw, kth, order=["size", "id"])[kth]
      )
  
      # test flatten case
      raw = np.random.rand(10, 10)
      x = tensor(raw, chunk_size=5)
  
      px = partition(x, kth, axis=None)
  
      res = px.execute().fetch()
      np.testing.assert_array_equal(
          res[kth], np.partition(raw, kth, axis=None)[kth]
      )
  
      # test multi-dimension
      raw = np.random.rand(10, 100)
      x = tensor(raw, chunk_size=(5, 20))
  
      kth = np.random.RandomState(0).randint(-10, 10, size=(3,))
      px = partition(x, kth)
  
      res = px.execute().fetch()
      np.testing.assert_array_equal(
          res[:, kth], np.partition(raw, kth)[:, kth]
      )
  
      raw = np.random.rand(10, 99)
      x = tensor(raw, chunk_size=(5, 20))
  
      px = partition(x, kth)
  
      res = px.execute().fetch()
      np.testing.assert_array_equal(
          res[:, kth], np.partition(raw, kth)[:, kth]
      )
  
      # test 3-d
      raw = np.random.rand(20, 25, 28)
      x = tensor(raw, chunk_size=(10, 15, 14))
  
      kth = np.random.RandomState(0).randint(-28, 28, size=(3,))
      px = partition(x, kth)
  
      res = px.execute().fetch()
      np.testing.assert_array_equal(
          res[:, :, kth], np.partition(raw, kth)[:, :, kth]
      )
  
      kth = np.random.RandomState(0).randint(-20, 20, size=(3,))
      px = partition(x, kth, axis=0)
  
      res = px.execute().fetch()
      np.testing.assert_array_equal(
          res[kth], np.partition(raw, kth, axis=0)[kth]
      )
  
      kth = np.random.RandomState(0).randint(-25, 25, size=(3,))
      px = partition(x, kth, axis=1)
  
      res = px.execute().fetch()
      np.testing.assert_array_equal(
          res[:, kth], np.partition(raw, kth, axis=1)[:, kth]
      )
  
      # test multi-dimension with structured type
      raw = np.empty(
          (10, 100), dtype=[("id", np.int32), ("size", np.int64)]
      )
      raw["id"] = np.random.randint(
          1000, size=(10, 100), dtype=np.int32
      )
      raw["size"] = np.random.randint(
          1000, size=(10, 100), dtype=np.int64
      )
      x = tensor(raw, chunk_size=(7, 30))
  
      kth = np.random.RandomState(0).randint(-100, 100, size=(10,))
      px = partition(x, kth)
  
      res = px.execute().fetch()
      np.testing.assert_array_equal(
          res[:, kth], np.partition(raw, kth)[:, kth]
      )
  
      px = partition(x, kth, order=["size", "id"])
  
      res = px.execute().fetch()
      np.testing.assert_array_equal(
          res[:, kth],
          np.partition(raw, kth, order=["size", "id"])[:, kth],
      )
  
      px = partition(x, kth, order=["size"])
  
      res = px.execute().fetch()
      np.testing.assert_array_equal(
          res[:, kth], np.partition(raw, kth, order=["size"])[:, kth]
      )
  
      kth = np.random.RandomState(0).randint(-10, 10, size=(5,))
      px = partition(x, kth, axis=0, order=["size", "id"])
  
      res = px.execute().fetch()
      np.testing.assert_array_equal(
          res[kth],
          np.partition(raw, kth, axis=0, order=["size", "id"])[kth],
      )
  
      raw = np.random.rand(10, 12)
      a = tensor(raw, chunk_size=(5, 4))
      kth = np.random.RandomState(0).randint(-12, 12, size=(2,))
      a.partition(kth, axis=1)
  
      res = a.execute().fetch()
      np.testing.assert_array_equal(
          res[:, kth], np.partition(raw, kth, axis=1)[:, kth]
      )
  
      kth = np.random.RandomState(0).randint(-10, 10, size=(2,))
      a.partition(kth, axis=0)
  
      raw_base = res
      res = a.execute().fetch()
      np.testing.assert_array_equal(
          res[kth], np.partition(raw_base, kth, axis=0)[kth]
      )
  
      # test kth which is tensor
      raw = np.random.rand(10, 12)
      a = tensor(raw, chunk_size=(3, 5))
      kth = (mt.random.rand(5) * 24 - 12).astype(int)
  
      px = partition(a, kth)
      sx = sort(a)
  
      res = px.execute().fetch()
      kth_res = kth.execute().fetch()
      sort_res = sx.execute().fetch()
      np.testing.assert_array_equal(
          res[:, kth_res], sort_res[:, kth_res]
      )
  
      a = tensor(raw, chunk_size=(10, 12))
      kth = (mt.random.rand(5) * 24 - 12).astype(int)
  
      px = partition(a, kth)
      sx = sort(a)
  
      res = px.execute().fetch()
      kth_res = kth.execute().fetch()
      sort_res = sx.execute().fetch()
      np.testing.assert_array_equal(
          res[:, kth_res], sort_res[:, kth_res]
      )
  
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 35 for calling a global function with example usage: test_partition_execution(setup) without return types