code: def train_test_split(
      X,
      y=None,
      test_size: Optional[Union[float, int]] = None,
      train_size: Optional[Union[float, int]] = None,
      shuffle: bool = True,
      random_state: Optional[
          Union[int, cp.random.RandomState, np.random.RandomState]
      ] = None,
      stratify=None,
  ):
      """
      Partitions device data into four collated objects, mimicking
      Scikit-learn's `train_test_split
      <https://scikit-learn.org/stable/modules/generated/sklearn.model_selection.train_test_split.html>`_.
  
      Parameters
      ----------
      X : cudf.DataFrame or cuda_array_interface compliant device array
          Data to split, has shape (n_samples, n_features)
      y : str, cudf.Series or cuda_array_interface compliant device array
          Set of labels for the data, either a series of shape (n_samples) or
          the string label of a column in X (if it is a cuDF DataFrame)
          containing the labels
      train_size : float or int, optional
          If float, represents the proportion [0, 1] of the data
          to be assigned to the training set. If an int, represents the number
          of instances to be assigned to the training set. Defaults to 0.8
      shuffle : bool, optional
          Whether or not to shuffle inputs before splitting
      random_state : int, CuPy RandomState or NumPy RandomState optional
          If shuffle is true, seeds the generator. Unseeded by default
  
      stratify: cudf.Series or cuda_array_interface compliant device array,
              optional parameter. When passed, the input is split using this
              as column to startify on. Default=None
  
      Examples
      --------
      .. code-block:: python
  
          >>> import cudf
          >>> from cuml.model_selection import train_test_split
          >>> # Generate some sample data
          >>> df = cudf.DataFrame({'x': range(10),
          ...                      'y': [0, 1] * 5})
          >>> print(f'Original data: {df.shape[0]} elements')
          Original data: 10 elements
          >>> # Suppose we want an 80/20 split
          >>> X_train, X_test, y_train, y_test = train_test_split(df, 'y',
          ...                                                     train_size=0.8)
          >>> print(f'X_train: {X_train.shape[0]} elements')
          X_train: 8 elements
          >>> print(f'X_test: {X_test.shape[0]} elements')
          X_test: 2 elements
          >>> print(f'y_train: {y_train.shape[0]} elements')
          y_train: 8 elements
          >>> print(f'y_test: {y_test.shape[0]} elements')
          y_test: 2 elements
  
          >>> # Alternatively, if our labels are stored separately
          >>> labels = df['y']
          >>> df = df.drop(['y'], axis=1)
          >>> # we can also do
          >>> X_train, X_test, y_train, y_test = train_test_split(df, labels,
          ...                                                     train_size=0.8)
  
      Returns
      -------
  
      X_train, X_test, y_train, y_test : cudf.DataFrame or array-like objects
          Partitioned dataframes if X and y were cuDF objects. If `y` was
          provided as a column name, the column was dropped from `X`.
          Partitioned numba device arrays if X and y were Numba device arrays.
          Partitioned CuPy arrays for any other input.
  
      """
      if isinstance(y, str):
          # Use the column with name `str` as y
          if isinstance(X, cudf.DataFrame):
              name = y
              y = X[name]
              X = X.drop(name, axis=1)
          else:
              raise TypeError(
                  "X needs to be a cuDF Dataframe when y is a \
                               string"
              )
  
      # todo: this check will be replaced with upcoming improvements
      # to input_utils
      #
      if y is not None:
          if not hasattr(
              X, "__cuda_array_interface__"
          ) and not isinstance(X, cudf.DataFrame):
              raise TypeError(
                  "X needs to be either a cuDF DataFrame, Series or \
                              a cuda_array_interface compliant array."
              )
  
          if not hasattr(
              y, "__cuda_array_interface__"
          ) and not isinstance(y, cudf.DataFrame):
              raise TypeError(
                  "y needs to be either a cuDF DataFrame, Series or \
                              a cuda_array_interface compliant array."
              )
  
          if X.shape[0] != y.shape[0]:
              raise ValueError(
                  "X and y must have the same first dimension"
                  "(found {} and {})".format(X.shape[0], y.shape[0])
              )
      else:
          if not hasattr(
              X, "__cuda_array_interface__"
          ) and not isinstance(X, cudf.DataFrame):
              raise TypeError(
                  "X needs to be either a cuDF DataFrame, Series or \
                              a cuda_array_interface compliant object."
              )
  
      if isinstance(train_size, float):
          if not 0 <= train_size <= 1:
              raise ValueError(
                  "proportion train_size should be between"
                  "0 and 1 (found {})".format(train_size)
              )
  
      if isinstance(train_size, int):
          if not 0 <= train_size <= X.shape[0]:
              raise ValueError(
                  "Number of instances train_size should be between 0 and the"
                  "first dimension of X (found {})".format(train_size)
              )
  
      if isinstance(test_size, float):
          if not 0 <= test_size <= 1:
              raise ValueError(
                  "proportion test_size should be between"
                  "0 and 1 (found {})".format(train_size)
              )
  
      if isinstance(test_size, int):
          if not 0 <= test_size <= X.shape[0]:
              raise ValueError(
                  "Number of instances test_size should be between 0 and the"
                  "first dimension of X (found {})".format(test_size)
              )
  
      x_numba = cuda.devicearray.is_cuda_ndarray(X)
      y_numba = cuda.devicearray.is_cuda_ndarray(y)
  
      # Determining sizes of splits
      if isinstance(train_size, float):
          train_size = int(X.shape[0] * train_size)
  
      if test_size is None:
          if train_size is None:
              train_size = int(X.shape[0] * 0.75)
  
          test_size = X.shape[0] - train_size
  
      if isinstance(test_size, float):
          test_size = int(X.shape[0] * test_size)
          if train_size is None:
              train_size = X.shape[0] - test_size
  
      elif isinstance(test_size, int):
          if train_size is None:
              train_size = X.shape[0] - test_size
  
      if shuffle:
          # Shuffle the data
          if random_state is None or isinstance(random_state, int):
              idxs = cp.arange(X.shape[0])
              random_state = cp.random.RandomState(seed=random_state)
  
          elif isinstance(random_state, cp.random.RandomState):
              idxs = cp.arange(X.shape[0])
  
          elif isinstance(random_state, np.random.RandomState):
              idxs = np.arange(X.shape[0])
  
          else:
              raise TypeError(
                  "`random_state` must be an int, NumPy RandomState \
                               or CuPy RandomState."
              )
  
          random_state.shuffle(idxs)
  
          if isinstance(X, cudf.DataFrame) or isinstance(
              X, cudf.Series
          ):
              X = X.iloc[idxs]
  
          elif hasattr(X, "__cuda_array_interface__"):
              # numba (and therefore rmm device_array) does not support
              # fancy indexing
              X = cp.asarray(X)[idxs]
  
          if isinstance(y, cudf.DataFrame) or isinstance(
              y, cudf.Series
          ):
              y = y.iloc[idxs]
  
          elif hasattr(y, "__cuda_array_interface__"):
              y = cp.asarray(y)[idxs]
  
          if stratify is not None:
              if isinstance(stratify, cudf.DataFrame) or isinstance(
                  stratify, cudf.Series
              ):
                  stratify = stratify.iloc[idxs]
  
              elif hasattr(stratify, "__cuda_array_interface__"):
                  stratify = cp.asarray(stratify)[idxs]
  
              split_return = _stratify_split(
                  X,
                  stratify,
                  y,
                  train_size,
                  test_size,
                  x_numba,
                  y_numba,
                  random_state,
              )
              return split_return
  
      # If not stratified, perform train_test_split splicing
      x_order = array_to_memory_order(X)
  
      if y is None:
          y_order = None
      else:
          y_order = array_to_memory_order(y)
  
      if hasattr(X, "__cuda_array_interface__") or isinstance(
          X, cupyx.scipy.sparse.csr_matrix
      ):
          X_train = cp.array(X[0:train_size], order=x_order)
          X_test = cp.array(X[-1 * test_size :], order=x_order)
          if y is not None:
              y_train = cp.array(y[0:train_size], order=y_order)
              y_test = cp.array(y[-1 * test_size :], order=y_order)
      elif isinstance(X, cudf.DataFrame):
          X_train = X.iloc[0:train_size]
          X_test = X.iloc[-1 * test_size :]
          if y is not None:
              if isinstance(y, cudf.Series):
                  y_train = y.iloc[0:train_size]
                  y_test = y.iloc[-1 * test_size :]
              elif hasattr(y, "__cuda_array_interface__") or isinstance(
                  y, cupyx.scipy.sparse.csr_matrix
              ):
                  y_train = cp.array(y[0:train_size], order=y_order)
                  y_test = cp.array(y[-1 * test_size :], order=y_order)
  
      if x_numba:
          X_train = cuda.as_cuda_array(X_train)
          X_test = cuda.as_cuda_array(X_test)
  
      if y_numba:
          y_train = cuda.as_cuda_array(y_train)
          y_test = cuda.as_cuda_array(y_test)
  
      if y is not None:
          return X_train, X_test, y_train, y_test
      else:
          return X_train, X_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 3 for calling a global function with example usage: train_test_split(X, y, test_size, train_size, shuffle, random_state, stratify) and returns: X_train,  X_test,  y_train,  y_test, X_train,  X_test, split_return