dask array shape is lost after using test_train_split from dask_ml

Question

In preparation for a machine learning task, I would like to perform a test/train split on an array that is stored in zarr. The array has shape (N, 288, 5) with