code: def test_rewrite_simple_shuffle_layer(ray_start_regular_shared):
      npartitions = 10
      df = dd.from_pandas(
          pd.DataFrame(
              np.random.randint(0, 100, size=(100, 2)),
              columns=["age", "grade"],
          ),
          npartitions=npartitions,
      )
      # We set max_branch=npartitions in order to ensure that the task-based
      # shuffle happens in a single stage, which is required in order for our
      # optimization to work.
      a = df.set_index(["age"], shuffle="tasks", max_branch=npartitions)
  
      dsk = a.__dask_graph__()
      keys = a.__dask_keys__()
      assert any(
          type(v) is SimpleShuffleLayer for k, v in dsk.layers.items()
      )
      dsk = rewrite_simple_shuffle_layer(dsk, keys)
      assert all(
          type(v) is not SimpleShuffleLayer
          for k, v in dsk.layers.items()
      )
      assert any(
          type(v) is MultipleReturnSimpleShuffleLayer
          for k, v in dsk.layers.items()
      )
  
file_path: ray/python/ray/util/dask/tests/test_dask_optimization.py
imports: import dask
  import dask.dataframe
  from dask.dataframe.shuffle import SimpleShuffleLayer
  from unittest import mock
  import numpy
  import pandas
  import pytest
  from ray.tests.conftest import *
  from ray.util.dask import dataframe_optimize
  from ray.util.dask.optimizations import rewrite_simple_shuffle_layer, MultipleReturnSimpleShuffleLayer
label_desc: Use image node_id 1 for calling a global function with example usage: test_rewrite_simple_shuffle_layer(ray_start_regular_shared) without return types