code: def _stratify_split(
      X,
      stratify,
      labels,
      n_train,
      n_test,
      x_numba,
      y_numba,
      random_state,
  ):
      """
      Function to perform a stratified split based on stratify column.
      Based on scikit-learn stratified split implementation.
  
      Parameters
      ----------
      X, y: Shuffled input data and labels
      stratify: column to be stratified on.
      n_train: Number of samples in train set
      n_test: number of samples in test set
      x_numba: Determines whether the data should be converted to numba
      y_numba: Determines whether the labales should be converted to numba
  
      Returns
      -------
      X_train, X_test: Data X divided into train and test sets
      y_train, y_test: Labels divided into train and test sets
      """
      x_cudf = False
      labels_cudf = False
  
      if isinstance(X, cudf.DataFrame):
          x_cudf = True
      elif hasattr(X, "__cuda_array_interface__"):
          X = cp.asarray(X)
      x_order = array_to_memory_order(X)
  
      # labels and stratify will be only cp arrays
      if isinstance(labels, cudf.Series):
          labels_cudf = True
          labels = labels.values
      elif hasattr(labels, "__cuda_array_interface__"):
          labels = cp.asarray(labels)
      elif isinstance(stratify, cudf.DataFrame):
          # ensuring it has just one column
          if labels.shape[1] != 1:
              raise ValueError(
                  "Expected one column for labels, but found df"
                  "with shape = %d" % (labels.shape)
              )
          labels_cudf = True
          labels = labels[0].values
  
      labels_order = array_to_memory_order(labels)
  
      # Converting to cupy array removes the need to add an if-else block
      # for startify column
      if isinstance(stratify, cudf.Series):
          stratify = stratify.values
      elif hasattr(stratify, "__cuda_array_interface__"):
          stratify = cp.asarray(stratify)
      elif isinstance(stratify, cudf.DataFrame):
          # ensuring it has just one column
          if stratify.shape[1] != 1:
              raise ValueError(
                  "Expected one column, but found column"
                  "with shape = %d" % (stratify.shape)
              )
          stratify = stratify[0].values
  
      classes, stratify_indices = cp.unique(
          stratify, return_inverse=True
      )
  
      n_classes = classes.shape[0]
      class_counts = cp.bincount(stratify_indices)
      if cp.min(class_counts) < 2:
          raise ValueError(
              "The least populated class in y has only 1"
              " member, which is too few. The minimum"
              " number of groups for any class cannot"
              " be less than 2."
          )
  
      if n_train < n_classes:
          raise ValueError(
              "The train_size = %d should be greater or "
              "equal to the number of classes = %d"
              % (n_train, n_classes)
          )
  
      class_indices = cp.split(
          cp.argsort(stratify_indices),
          cp.cumsum(class_counts)[:-1].tolist(),
      )
  
      X_train = None
  
      # random_state won't be None or int, that's handled earlier
      if isinstance(random_state, np.random.RandomState):
          random_state = cp.random.RandomState(
              seed=random_state.get_state()[1]
          )
  
      # Break ties
      n_i = _approximate_mode(class_counts, n_train, random_state)
      class_counts_remaining = class_counts - n_i
      t_i = _approximate_mode(
          class_counts_remaining, n_test, random_state
      )
  
      for i in range(n_classes):
          permutation = random_state.permutation(class_counts[i].item())
          perm_indices_class_i = class_indices[i].take(permutation)
  
          y_train_i = cp.array(
              labels[perm_indices_class_i[: n_i[i]]], order=labels_order
          )
          y_test_i = cp.array(
              labels[perm_indices_class_i[n_i[i] : n_i[i] + t_i[i]]],
              order=labels_order,
          )
          if hasattr(X, "__cuda_array_interface__") or isinstance(
              X, cupyx.scipy.sparse.csr_matrix
          ):
              X_train_i = cp.array(
                  X[perm_indices_class_i[: n_i[i]]], order=x_order
              )
              X_test_i = cp.array(
                  X[perm_indices_class_i[n_i[i] : n_i[i] + t_i[i]]],
                  order=x_order,
              )
  
              if X_train is None:
                  X_train = cp.array(X_train_i, order=x_order)
                  y_train = cp.array(y_train_i, order=labels_order)
                  X_test = cp.array(X_test_i, order=x_order)
                  y_test = cp.array(y_test_i, order=labels_order)
              else:
                  X_train = cp.concatenate([X_train, X_train_i], axis=0)
                  X_test = cp.concatenate([X_test, X_test_i], axis=0)
                  y_train = cp.concatenate([y_train, y_train_i], axis=0)
                  y_test = cp.concatenate([y_test, y_test_i], axis=0)
  
          elif x_cudf:
              X_train_i = X.iloc[perm_indices_class_i[: n_i[i]]]
              X_test_i = X.iloc[
                  perm_indices_class_i[n_i[i] : n_i[i] + t_i[i]]
              ]
  
              if X_train is None:
                  X_train = X_train_i
                  y_train = y_train_i
                  X_test = X_test_i
                  y_test = y_test_i
              else:
                  X_train = cudf.concat(
                      [X_train, X_train_i], ignore_index=False
                  )
                  X_test = cudf.concat(
                      [X_test, X_test_i], ignore_index=False
                  )
                  y_train = cp.concatenate([y_train, y_train_i], axis=0)
                  y_test = cp.concatenate([y_test, y_test_i], axis=0)
  
      if x_numba:
          X_train = cuda.as_cuda_array(X_train)
          X_test = cuda.as_cuda_array(X_test)
      elif x_cudf:
          X_train = cudf.DataFrame(X_train)
          X_test = cudf.DataFrame(X_test)
  
      if y_numba:
          y_train = cuda.as_cuda_array(y_train)
          y_test = cuda.as_cuda_array(y_test)
      elif labels_cudf:
          y_train = cudf.Series(y_train)
          y_test = cudf.Series(y_test)
  
      return X_train, X_test, y_train, y_test
  
file_path: cuml/python/cuml/model_selection/_split.py
imports: from typing import Optional, Union
  from cuml.common import input_to_cuml_array
  from cuml.internals.array import array_to_memory_order
  from cuml.internals.safe_imports import cpu_only_import, gpu_only_import, gpu_only_import_from
label_desc: Use image node_id 1 for calling a global function with example usage: _stratify_split(X, stratify, labels, n_train, n_test, x_numba, y_numba, random_state) and returns: X_train,  X_test,  y_train,  y_test