code: def test_shuffle_execution(setup):
      # test consistency
      s1 = np.arange(9).reshape(3, 3)
      s2 = np.arange(1, 10).reshape(3, 3)
      ts1 = mt.array(s1, chunk_size=2)
      ts2 = mt.array(s2, chunk_size=2)
  
      ret = shuffle(ts1, ts2, axes=[0, 1], random_state=0)
      res1, res2 = ret.execute().fetch()
  
      # calc row index
      s1_col_0 = s1[:, 0].tolist()
      rs1_col_0 = [
          res1[:, i]
          for i in range(3)
          if set(s1_col_0) == set(res1[:, i])
      ][0]
      row_index = [s1_col_0.index(j) for j in rs1_col_0]
      # calc col index
      s1_row_0 = s1[0].tolist()
      rs1_row_0 = [
          res1[i] for i in range(3) if set(s1_row_0) == set(res1[i])
      ][0]
      col_index = [s1_row_0.index(j) for j in rs1_row_0]
      np.testing.assert_array_equal(res2, s2[row_index][:, col_index])
  
      # tensor + tensor
      raw1 = np.random.rand(10, 15, 20)
      t1 = mt.array(raw1, chunk_size=8)
      raw2 = np.random.rand(10, 15, 20)
      t2 = mt.array(raw2, chunk_size=5)
  
      for axes in [(0,), (0, 1), (0, 2), (1, 2), (0, 1, 2)]:
          ret = shuffle(t1, t2, axes=axes, random_state=0)
          res1, res2 = ret.execute().fetch()
  
          assert res1.shape == raw1.shape
          assert res2.shape == raw2.shape
          np.testing.assert_array_equal(
              _sort(raw1, axes), _sort(res1, axes)
          )
          np.testing.assert_array_equal(
              _sort(raw2, axes), _sort(res2, axes)
          )
  
      # tensor + tensor(more dimension)
      raw3 = np.random.rand(10, 15)
      t3 = mt.array(raw3, chunk_size=(8, 15))
      raw4 = np.random.rand(10, 15, 20)
      t4 = mt.array(raw4, chunk_size=(5, 15, 10))
  
      for axes in [(1,), (0, 1), (1, 2)]:
          ret = shuffle(t3, t4, axes=axes, random_state=0)
          res3, res4 = ret.execute().fetch()
  
          assert res3.shape == raw3.shape
          assert res4.shape == raw4.shape
          np.testing.assert_array_equal(
              _sort(raw3, axes), _sort(res3, axes)
          )
          np.testing.assert_array_equal(
              _sort(raw4, axes), _sort(res4, axes)
          )
  
      # tensor + dataframe + series
      raw5 = np.random.rand(10, 15, 20)
      t5 = mt.array(raw5, chunk_size=8)
      t6 = mt.array(raw5[:, 0, 0], chunk_size=6)
      raw6 = pd.DataFrame(np.random.rand(10, 15))
      df = md.DataFrame(raw6, chunk_size=(8, 15))
      raw7 = pd.Series(np.random.rand(10))
      series = md.Series(raw7, chunk_size=8)
  
      for axes in [(0,), (1,), (0, 1), (1, 2), [0, 1, 2]]:
          ret = shuffle(t5, df, series, t6, axes=axes, random_state=0)
          # skip check nsplits because it's updated
          res5, res_df, res_series, res6 = ret.execute(
              extra_config={"check_nsplits": False}
          ).fetch(extra_config={"check_nsplits": False})
  
          assert res5.shape == raw5.shape
          assert res_df.shape == df.shape
          assert res_series.shape == series.shape
          assert res6.shape == (raw5.shape[0],)
  
file_path: mars/mars/learn/utils/tests/test_shuffle.py
imports: import numpy
  import pandas
  import pytest
  from ....None import tensor
  from ....None import dataframe
  from ....core import tile
  from ..None import shuffle
  from ..shuffle import LearnShuffle
label_desc: Use image node_id 3 for calling a global function with example usage: test_shuffle_execution(setup) without return types