code: def test_extract_partitions_shape(
      nrows, ncols, n_parts, input_type, colocated, client
  ):
      adj_input_type = (
          "dataframe" if input_type == "series" else input_type
      )
  
      X_arr, y_arr = make_blobs(
          n_samples=nrows, n_features=ncols, n_parts=n_parts
      )
  
      if adj_input_type == "dataframe" or input_type == "dataframe":
          X = to_dask_cudf(X_arr)
          y = to_dask_cudf(y_arr)
      elif input_type == "array":
          X, y = X_arr, y_arr
  
      if input_type == "series":
          X = X[X.columns[0]]
  
      if input_type == "dataframe" or input_type == "series":
          X_len_parts = X.map_partitions(len).compute()
          y_len_parts = y.map_partitions(len).compute()
      elif input_type == "array":
          X_len_parts = X.chunks[0]
          y_len_parts = y.chunks[0]
  
      if colocated:
          ddh = DistributedDataHandler.create((X, y), client)
          parts = [part.result() for worker, part in ddh.gpu_futures]
          for i in range(len(parts)):
              assert (parts[i][0].shape[0] == X_len_parts[i]) and (
                  parts[i][1].shape[0] == y_len_parts[i]
              )
      else:
          ddh = DistributedDataHandler.create(X, client)
          parts = [part.result() for worker, part in ddh.gpu_futures]
          for i in range(len(parts)):
              assert parts[i].shape[0] == X_len_parts[i]
  
file_path: cuml/python/cuml/tests/dask/test_dask_input_utils.py
imports: from cuml.dask.common.dask_arr_utils import to_dask_cudf
  import pytest
  from cuml.dask.datasets.blobs import make_blobs
  from cuml.dask.common.input_utils import DistributedDataHandler
  import dask.array
  from cuml.internals.safe_imports import gpu_only_import
label_desc: Use image node_id 2 for calling a global function with example usage: test_extract_partitions_shape(nrows, ncols, n_parts, input_type, colocated, client) without return types