train_test_split fails for too many values (32bit only)

Consider the following code:

import numpy as np
from sklearn.model_selection import train_test_split

n = 10000
y = np.random.randint(0, 2, size=n)

y_train, y_test = train_test_split(y, train_size=int(n/2),
                                   test_size=int(n/2), stratify=y, random_state=123)

print('num train: {}'.format(len(y_train)))
print('train mean: {}'.format(y_train.mean()))
print('num test: {}'.format(len(y_test)))
print('test mean: {}'.format(y_test.mean()))

When n=10,000, I correctly obtain:

num train: 5000
train mean: 0.4958
num test: 5000
test mean: 0.4958

But for larger n, such as n=100,000, I get the following error:

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
/home/scott/Development/scratch/sklearn/stratify.py in <module>()
      6 y = np.random.randint(0, 2, size=n)
      7 
----> 8 y_train, y_test = train_test_split(y, train_size=n/2, test_size=n/2, stratify=y, random_state=123)
      9 
     10 print 'num train: {}'.format(len(y_train))

/home/scott/anaconda/lib/python2.7/site-packages/sklearn/model_selection/_split.pyc in train_test_split(*arrays, **options)
   1700     train, test = next(cv.split(X=arrays[0], y=stratify))
   1701     return list(chain.from_iterable((safe_indexing(a, train),
-> 1702                                      safe_indexing(a, test)) for a in arrays))
   1703 
   1704 

/home/scott/anaconda/lib/python2.7/site-packages/sklearn/model_selection/_split.pyc in <genexpr>((a,))
   1700     train, test = next(cv.split(X=arrays[0], y=stratify))
   1701     return list(chain.from_iterable((safe_indexing(a, train),
-> 1702                                      safe_indexing(a, test)) for a in arrays))
   1703 
   1704 

/home/scott/anaconda/lib/python2.7/site-packages/sklearn/utils/__init__.pyc in safe_indexing(X, indices)
    110             return X.take(indices, axis=0)
    111         else:
--> 112             return X[indices]
    113     else:
    114         return [X[idx] for idx in indices]

IndexError: arrays used as indices must be of integer (or boolean) type

And for n=1,000,000, I don’t get an exception, instead the strange results:

num train: 1785
train mean: 0.414565826331
num test: 894
test mean: 0.414988814318

Why is this? Is this a bug? Does train_test_split fail with too many values?

[aqua4]

Hello! Your code returns correct results without any errors on my system. I’m using: python 2.7.12 sklearn 0.18.1 numpy 1.11.3

[amueller]

Could it be an issue in

        class_indices = np.split(np.argsort(y_indices, kind='mergesort'),
                                 np.cumsum(class_counts)[:-1])

maybe converting class_counts to float will fix it? (this code is no fun)