code: def test_stratify_any_input(test_size, train_size):
      X = cudf.DataFrame({"x": range(10)})
      X["test_col"] = cudf.Series([10, 0, 0, 10, 10, 10, 0, 0, 10, 10])
      y = cudf.Series(([0] * (10 // 2)) + ([1] * (10 // 2)))
  
      X_train, X_test, y_train, y_test = train_test_split(
          X,
          y,
          train_size=train_size,
          test_size=test_size,
          shuffle=True,
          stratify=X["test_col"],
          random_state=15,
      )
      assert (X_train["x"].to_numpy() == X_train.index.to_numpy()).all()
      assert (X_test["x"].to_numpy() == X_test.index.to_numpy()).all()
  
      if train_size is not None:
          assert X_train.shape[0] == (int)(X.shape[0] * train_size)
  
      elif test_size is not None:
          assert X_test.shape[0] == (int)(X.shape[0] * test_size)
  
imports: from cuml.internals.safe_imports import gpu_only_import_from
  from cuml.datasets import make_classification
  from cuml.model_selection import train_test_split
  import pytest
  from cuml.internals.safe_imports import cpu_only_import
  from cuml.internals.safe_imports import gpu_only_import
label_desc: Use image node_id 15 for calling a global function with example usage: test_stratify_any_input(test_size, train_size) without return types