I think that even a conservative change like the following would be welcome:

chunk = x[slc].copy()

In Dask at least the copy is free for us.

We could also maybe send in a cupy.random.RandomState object rather than a numpy.random.RandomState object so that we generate the random values on the GPU instead of on the CPU. This doesn't solve the np.arange problem, but maybe that points us in another possible direction. We would still need to remove explicit np.asarray calls on inputs, and allow array-like arguments through in some cases.

This came up in conversation with @ogrisel a couple of months ago.

Also cc @jakirkham , who has been looking at related problems recently.

Yeah with np.arange, I was hoping we could write something like this...

np.cumsum(np.ones_like(X, dtype="i8", shape=(n_samples,))) - 1

There's also a call to np.full. Currently that looks like this...

np.full(n_splits, n_samples // n_splits, dtype=np.int)

We might want to generalize it to something like this...

np.full_like(X, n_samples // n_splits, shape=(n_splits,), dtype=np.int)

Right, so that opens up a question for Scikit-Learn maintainers of if they're willing to accept code like this that is less intuitive in a few places for the case where the inputs are Numpy-like arrays.

My guess is that accepting a user-defined RandomState object is fine (it looks like there is good precedent for this already) while replacing np.arange calls with ones_like/cumsum is maybe less welcome.

There's of course the point that this syntax only works with NumPy 1.17, which may be it's own kind of deal breaker. Anyways figured I'd float these snippets to start the conversation.

Also there's a separate discussion about having a way to cast between different arrays in a subthread of issue ( numpy/numpy#13831 ), which is ostensibly about handling asarray without forcing array-like things to become NumPy arrays.

cc @pentschev (whose involved in that discussion and also thinking about these things)

Yeah it'll be a while till we can use numpy 1.17.

My guess is that accepting a user-defined RandomState object is fine...

Yeah it would be great if RandomState created the array and shuffled it for us, but wasn't finding the API we would need to do that. .shuffle does not seem to work if one does rs.shuffle(5) for example.

Another concern that Ben helped identify earlier is that if we want to allow a user-defined RandomState object, we'd want to make sure this passes check_random_state. Currently that means the user-defined RandomState object would need to subclass NumPy's RandomState object (please see code below). This is a bit undesirable as we will end up carrying around a Cython base class that has some irrelevant state. Not sure how well this will behave in practice either. Would it be possible to generalize this check to merely make sure the RandomState satisfies some abstract base class or has some particular methods?

@jakirkham I don't see why not.

Would it be possible to generalize this check to merely make sure the RandomState satisfies some abstract base class or has some particular methods?

Which interface do you have in mind? The interface for cupy.random.RandomState and np.random.RandomState are pretty extensive.

we kinda only need to distinguish it from an integer or None, right? We could check for a single common method existing. If we then try to use some obscure method and it's not implemented, well that's a reasonable error.
If it has a uniform method it's much more likely to be something that implements sampling than an integer, I think ;)

cc @ogrisel (as we discussed this during EuroSciPy)

It would be great if someone wanted to try to prototype Array API spec support in scikit-learn Cross-Validation tools in a draft PR on top of the #22554 PR.