sklearn train_test_split on pandas stratify by multiple columns

Question

I\'m a relatively new user to sklearn and have run into some unexpected behavior in train_test_split from sklearn.model_selection. I have a pandas dataframe that I would like to

Answer 1

What version of scikit-learn are you using ? You can use sklearn.__version__ to check.

The prior to version 0.19.0, scikit-learn does not handle 2-dimensional stratification correctly. It is patched in 0.19.0.

It is describled in issue #9044.

Update your scikit-learn should fix the problem. If you can't update your scikit-learn, see this commit history here for the fix.