[BUG] `KFold` with `shuffle=True` returns training indices in a different order from scikit-learn

[apiqwe]

Describe the bug

cuml.model_selection.KFold returns training indices in a different order from the equivalent sklearn.model_selection.KFold call when shuffle=True and the same random_state is used.

For the deterministic three-sample example below, the first split returned by cuML is:

[2 1] [0]

while scikit-learn returns:

[1 2] [0]

The train/test membership is the same in both implementations; the discrepancy is specifically the ordering of the training indices. The remaining two folds match exactly.

This difference can affect code that expects cuML's scikit-learn-compatible splitter API to produce the same deterministic index ordering for the same input, n_splits, shuffle, and random_state.

Steps/Code to reproduce bug

cuML reproducer:

import numpy as np
from cuml.model_selection import KFold

X = np.array([
    [1, 2],
    [3, 4],
    [5, 6],
])

for train, test in KFold(
    n_splits=3,
    shuffle=True,
    random_state=1,
).split(X):
    print(train, test)

Output:

[2 1] [0]
[0 1] [2]
[0 2] [1]

For comparison, the equivalent scikit-learn code:

import numpy as np
from sklearn.model_selection import KFold

X = np.array([
    [1, 2],
    [3, 4],
    [5, 6],
])

for train, test in KFold(
    n_splits=3,
    shuffle=True,
    random_state=1,
).split(X):
    print(train, test)

Output:

[1 2] [0]
[0 1] [2]
[0 2] [1]

Expected behavior

For the same input and deterministic parameters, cuml.model_selection.KFold should return training and test indices in an ordering compatible with scikit-learn.

For this reproducer, the expected output is:

[1 2] [0]
[0 1] [2]
[0 2] [1]

In particular, the first training split should be returned as [1, 2] rather than [2, 1].

If cuML intentionally does not guarantee the ordering of indices within each train/test split, that difference should be documented clearly because the returned index arrays are observably different even though the fold membership is identical.

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:

conda list
# packages in environment at /opt/conda/envs/rapids-26.08:
#
# Name                              Version          Build                                         Channel
# 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

The discrepancy is limited to the order of indices within the first training split:

cuML:         train=[2, 1], test=[0]
scikit-learn: train=[1, 2], test=[0]

As sets, both training splits contain exactly the same samples {1, 2}. Therefore, this reproducer does not show an incorrect fold assignment; it shows an ordering incompatibility in the index arrays returned by split().

With shuffle=True, shuffling should determine which samples are assigned to each fold. Once the test fold is selected, the returned training indices should remain deterministic and compatible with scikit-learn for the same random seed. Here, cuML appears to preserve a shuffled ordering for the first training split rather than the ordering returned by scikit-learn.

This distinction matters for downstream code that directly compares split indices, serializes splits, uses the returned order to index data without an additional sort, or expects drop-in compatibility with scikit-learn's deterministic output.

[betatim]

I am not sure if this is a bug or not. Especially as the split of what is in the train and what is in the test set match scikit-learn. The order of samples does not matter to the algorithms that scikit-learn and cuml offer.

For me this means that this is not a bug. It might be a good idea to record this difference in the documentation though.

Do you have a specific use-case in mind where the order of the samples in a set matters?

[SIDDARTHAREDDY8]

@betatim — sharing the concrete cases where the ordering bites (I am the author of the paused PR #8633; leaving that PR alone per your note and posting here since the decision is being made on this issue).

1. cuML-vs-sklearn parity checks. The reason teams adopt cuml.model_selection is drop-in sklearn compatibility, and the standard migration check is assert_array_equal (or similar) between the two libraries outputs — including the split indices. Membership matching but order differing fails those harnesses, which is exactly how this issue surfaced.

2. X[train] row order feeds training. [2, 1] vs [1, 2] produces a differently-ordered training matrix. For order-sensitive training — SGD without reshuffling (different minibatch composition/order), partial_fit online learning — that changes the fitted model, so the difference is observably more than cosmetic.

3. It is a well-defined target, not an implementation detail. sklearn KFold returns indices sorted within each split; I verified this empirically across 144 (n_samples, n_splits, random_state) combinations against sklearn 1.9.1 — sorted output matched in every case. So "match sklearn" is deterministic and testable.

Documenting the difference (your suggestion) is a fair fallback, but matching keeps the API-compat promise intact, and the change itself is tiny — sorting both index arrays at yield, no API change, negligible cost.

Happy to go whichever way triage decides.

[csadorf]

I don’t think this represents a bug under cuML’s current compatibility guarantees. Our documentation says that we match familiar APIs where possible while allowing differences arising from GPU-specific implementations and optimizations. A fixed random_state should guarantee reproducibility within cuML, but it does not guarantee equivalence with scikit-learn’s implementation. This is true even for cuml.accel: its compatibility documentation promises comparable model quality, not numerically identical CPU and GPU results.

I concede that matching scikit-learn’s ordering could be beneficial in some circumstances, even here. I would consider changing the behavior to improve compatibility here and elsewhere if it can be done easily without a performance impact, particularly if the current results are not equivalent under a reasonable quality or semantic comparison.

Here, however, the train/test membership is identical (which wouldn't even necessarily need to be guaranteed IMO); only the ordering differs and the performance impact for matching is maybe not large, but certainly not neutral.

I would not accept this as a bug, but, if at all, as an optional compatibility improvement. Any "fix" would need to be performance neutral.

[SIDDARTHAREDDY8]

@csadorf - thanks for the triage decision, that's clear. I've reworked #8633 accordingly: it no longer changes default behavior and instead adds an opt-in sort_indices parameter to KFold (default False), i.e. the optional, performance-neutral compatibility improvement you described. Default output is verified byte-identical to current upstream behavior with no extra kernels. Details and verification numbers are in the PR body.