add params ensure_2d and allow_nd to check_array (#722)

* add params ensure_2d and allow_nd to check_array

allow to use more array with more than 2 dimensions

Author: sim-san

dask_ml/model_selection/_split.py

@@ -148,7 +148,7 @@ def __init__(
         self.blockwise = _check_blockwise(blockwise)
 
     def split(self, X, y=None, groups=None):
-        X = check_array(X)
+        X = check_array(X, ensure_2d=False, allow_nd=True)
         rng = check_random_state(self.random_state)
         for i in range(self.n_splits):
             seeds = draw_seed(rng, 0, _I4MAX, size=len(X.chunks[0]), dtype="uint")

tests/model_selection/test_split.py

@@ -240,3 +240,14 @@ def test_split_mixed():
     assert len(expected) == len(results)
     for a, b in zip(expected, results):
         da.utils.assert_eq(a, b)
+
+
+def test_split_3d_data():
+    X_3d = np.arange(1.0, 5001.0).reshape((100, 10, 5))
+    y_3d = np.arange(1.0, 101.0).reshape(100, 1)
+
+    r = dask_ml.model_selection.train_test_split(X_3d, y_3d)
+    X_train, X_test, y_train, y_test = r
+
+    assert X_train.ndim == X_3d.ndim
+    assert X_train.shape[1:] == X_3d.shape[1:]