def test_reshape_all_not_chunked_merge(inshape, inchunks, expected_inchunks, outshape, outchunks): base = np.arange(np.prod(inshape)).reshape(inshape) a = da.from_array(base, chunks=inchunks) # test directly inchunks2, outchunks2 = reshape_rechunk(a.shape, outshape, inchunks) assert inchunks2 == expected_inchunks assert outchunks2 == outchunks # and via reshape result = a.reshape(outshape) assert result.chunks == outchunks assert_eq(result, base.reshape(outshape))
def test_reshape_all_chunked_no_merge(inshape, inchunks, outshape, outchunks): # https://github.com/dask/dask/issues/5544#issuecomment-712280433 # When the early axes are completely chunked then we are just moving blocks # and can avoid any rechunking. The result inchunks are the same as the # input chunks. base = np.arange(np.prod(inshape)).reshape(inshape) a = da.from_array(base, chunks=inchunks) # test directly inchunks2, outchunks2 = reshape_rechunk(a.shape, outshape, inchunks) assert inchunks2 == inchunks assert outchunks2 == outchunks # and via reshape result = a.reshape(outshape) assert result.chunks == outchunks assert_eq(result, base.reshape(outshape))
def test_reshape_rechunk(inshape, outshape, prechunks, inchunks, outchunks): result_in, result_out = reshape_rechunk(inshape, outshape, prechunks) assert result_in == inchunks assert result_out == outchunks assert np.prod(list(map(len, result_in))) == np.prod(list(map(len, result_out)))