I'm not sure how to make tests for the stratify and shuffle cases. Maybe the example I added for stratify is enough? Other than that I suppose this PR is ready.

Ideally the shuffling should not alter too much the balance in the folds. Otherwise, what is the difference with GroupShuffleSplit?

I modified "stratify" to use the median value for each group.

The thing about balancing the folds vs. stratifying or shuffling is that it is a trade-off, and this PR allows one to make the choice explicitly. If your dataset has a small number of groups, they will strongly constrain the possible folds, and balancing might have priority. When there is a sufficient number of groups, or the groups are all the same size, the sizes of folds will be approximately balanced either way, and one can stratify or shuffle the folds instead.

Tests appear to be failing

All tests pass now, so I went ahead and renamed this MRG.

It's not clear from the docs what this actually does. Can you maybe give an example in the docs?
So it computes the median y value in each group. What does that mean for categories? Wouldn't you want the mode? And for classification this would mostly have an impact if the folds have very different y distributions?

Can you please add a legend to the plots?
And would it be possible to show an effect with 5-fold CV?

And shuffle doesn't seem to assign groups randomly, but greedily in random order.
The description of balanced seems wrong, too. To do balanced you'd have to solve something like a bin-packing problem, right? We're only using a greedy heuristic.

Thanks for the comments. You're right, I didn't think the discrete case through, and the documentation should be improved.

What kind of API should I use for distinguishing the discrete and the continuous case? Autodetection is possible but probably not a good idea. I see that in the feature_selection module there are separate regression and classification functions; but since the code will be very similar, maybe a parameter is best?

The regular StratifiedKFold only supports the discrete case. Maybe I should add support for the continuous case there too?

You're right shuffling is not completely random, but otherwise the folds could be very uneven. Did you mean it should be better documented, or that the implementation should be different?
I understand that the balanced option does not give an optimal solution, I didn't mean to imply that.

My comments were mostly on the documentation. I think we should be very clear on what we are doing. Your solution seems good for regression, but is not the only possible way to stratify, I think.
I'm not opposed to adding this stratification strategy, but we should try to describe well what the different strategies are doing.

For feature_selection I think doing regression and classification in the same class has given us a lot of headache. So in some way I'd prefer different classes for regression and classification. On the other hand I like the way that stratification is implemented as an option here - it's somewhat different to the other classes, but we have a real explosion of different classes, and I might prefer to have it as a parameter.
But we can't really have different classes for regression and classification and have stratification as a parameter, that makes no sense. We could auto-detect using type_of_target, though that's a bit dangerous. If there are many different ints, how do you know if they are classes or regression targets?

maybe method="stratify_classification" and method="stratify_regression" would work so we don't have redundant parameters?

There is a PR somewhere for stratified cross-validation for regression. It uses binning, though, and I think sorting is better. I haven't looked at it in a while.

I think both are important.

But yeah, definitely we need a motivation for each case. I think the case for having some stratified option is pretty clear, but I feel there are non-obvious choices to make.

Closing this one as it needs a refresh and motivation, but related: #26821