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)))
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)))
