Skip to content

[BUG] StratifiedKFold rejects single-class targets that scikit-learn splits successfully #8632

Description

@apiqwe

Describe the bug

cuml.model_selection.StratifiedKFold rejects a target array containing only a single class, even though the equivalent sklearn.model_selection.StratifiedKFold call completes successfully.

For the following input:

X = np.arange(8).reshape(-1, 1)
y = np.zeros(8, dtype=int)

scikit-learn produces two valid folds:

[(array([4, 5, 6, 7]), array([0, 1, 2, 3])),
 (array([0, 1, 2, 3]), array([4, 5, 6, 7]))]

cuML instead raises:

ValueError: number of unique classes cannot be less than 2

The input size is sufficient for n_splits=2, and each fold can preserve the only available class distribution exactly, so the additional requirement for at least two unique classes prevents a case that scikit-learn handles normally.

Steps/Code to reproduce bug

cuML reproducer:

import numpy as np
from cuml.model_selection import StratifiedKFold

X = np.arange(8).reshape(-1, 1)
y = np.zeros(8, dtype=int)

print(list(StratifiedKFold(n_splits=2).split(X, y)))

Output:

Traceback (most recent call last):
  File "/workspace/apibughub/cuml/StratifiedKFold/error_bug1/error_bug_cuml.py", line 7, in <module>
    print(list(StratifiedKFold(n_splits=2).split(X, y)))
          ~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/envs/rapids-26.08/lib/python3.14/site-packages/cuml/model_selection/_split.py", line 185, in split
    yield from self._split(X, y, indices)
  File "/opt/conda/envs/rapids-26.08/lib/python3.14/site-packages/cuml/model_selection/_split.py", line 320, in _split
    raise ValueError("number of unique classes cannot be less than 2")
ValueError: number of unique classes cannot be less than 2

For comparison, the equivalent scikit-learn code:

import numpy as np
from sklearn.model_selection import StratifiedKFold

X = np.arange(8).reshape(-1, 1)
y = np.zeros(8, dtype=int)

print(list(StratifiedKFold(n_splits=2).split(X, y)))

Output:

[(array([4, 5, 6, 7]), array([0, 1, 2, 3])),
 (array([0, 1, 2, 3]), array([4, 5, 6, 7]))]

Expected behavior

cuml.model_selection.StratifiedKFold should allow a single-class target when the number of samples is sufficient for the requested number of folds.

With all eight samples belonging to the same class and n_splits=2, each test fold can contain four samples and therefore preserve the class distribution exactly.

For this reproducer, cuML should return two valid train/test splits equivalent to:

[(array([4, 5, 6, 7]), array([0, 1, 2, 3])),
 (array([0, 1, 2, 3]), array([4, 5, 6, 7]))]

The exact ordering of indices is less important than allowing the split and producing valid non-overlapping folds with the same membership semantics.

If single-class targets are intentionally unsupported, this limitation should be documented explicitly because the equivalent scikit-learn estimator accepts this input.

Environment details (please complete the following information):

  • Environment location: Docker
  • Linux Distro/Architecture: Ubuntu 24.04 / x86_64
  • GPU Model/Driver: NVIDIA GeForce RTX 4090 / 595.71.05
  • CUDA: 13.2
  • Method of cuDF & cuML install: conda

conda list:

# packages in environment at /opt/conda/envs/rapids-26.08:
#
# Name              Version       Build                                      Channel
python              3.14.6        h242f9ac_102_cp314                         conda-forge
numpy               2.4.6         py314h2b28147_0                            conda-forge
scipy               1.16.3        py314hf07bd8e_2                            conda-forge
scikit-learn        1.9.0         np2py314hf09ca88_0                         conda-forge
rapids              26.08.00      cuda13_260806_c2656556                     rapidsai
cuml                26.08.00      cuda13_cp311_abi3_260805_265b9da6          rapidsai
libcuml             26.08.00      cuda13_260805_265b9da6                     rapidsai
cudf                26.08.00      cuda13_cp311_abi3_260805_ff5b362d          rapidsai
libraft             26.08.00      cuda13_260805_ebf92684                     rapidsai
libraft-headers     26.08.00      cuda13_260805_ebf92684                     rapidsai
pylibraft           26.08.00      cuda13_cp311_abi3_260805_ebf92684          rapidsai
cuvs                26.08.01      cuda13_cp311_abi3_260806_25b1be43          rapidsai
libcuvs             26.08.01      cuda13_260806_25b1be43                     rapidsai
cupy                14.1.1        py314hdea9c46_0                            conda-forge
cupy-core           14.1.1        py314hcd3b49b_0                            conda-forge
numba               0.64.0        py314h8169c2f_0                            conda-forge
numba-cuda          0.30.4        py314h42812f9_0                            conda-forge
rmm                 26.08.00      cuda13_cp311_abi3_260805_42d059f1          rapidsai
librmm              26.08.00      cuda13_260805_42d059f1                     rapidsai
cuda-version        13.3           hcbadf70_3                                 conda-forge
cuda-bindings       13.3.1        py314h42812f9_1                            conda-forge
cuda-cudart         13.3.29       hecca717_0                                 conda-forge
cuda-nvrtc          13.3.33       hecca717_0                                 conda-forge
libcublas           13.6.0.2      h676940d_0                                 conda-forge
libcusolver         12.2.6.9      h676940d_0                                 conda-forge
libcusparse         12.8.2.51     hecca717_0                                 conda-forge
libcurand           10.4.3.29     h676940d_0                                 conda-forge

Additional context

This reproducer uses the simplest possible target distribution:

y = [0, 0, 0, 0, 0, 0, 0, 0]

Although there is only one class, stratification is still well-defined because every subset has the same class proportion: 100% class 0.

With n_splits=2, scikit-learn divides the eight samples into two four-sample test folds and returns complementary four-sample training folds.

The cuML implementation instead performs an explicit unique-class-count check and raises before creating any folds:

ValueError: number of unique classes cannot be less than 2

Therefore, the discrepancy appears to come from an extra input restriction in cuML rather than from an inability to construct the requested split.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions