Repository: dask/dask
File: dask/array/reshape.py (Python)

Code before:
def _calc_lower_dimension_chunks(chunks, start, stop):
    # We need the lower dimension chunks to match what the higher dimension chunks
    # can be combined to, i.e. multiply the different dimensions
    return tuple(
        map(
            lambda x: reduce(mul, x),
            product(*chunks[start : stop + 1]),
        )
    )


def _smooth_chunks(ileft, ii, max_in_chunk, result_inchunks):
    # The previous step squashed the whole dimension into a single
    # chunk for ileft + 1 (and potentially combined too many elements
    # into a single chunk for ileft as well). We split up the single
    # chunk into multiple chunks to match the max_in_chunk to keep
    # chunksizes consistent:
    # ((1, 1), (200)) -> ((1, 1), (20, ) * 10) for max_in_chunk = 20
    # It's important to ensure that all dimensions before the dimension
    # we adjust have all-1 chunks to respect C contiguous arrays
    # during the reshaping
    # Example:
    # Assume arr = da.from_array(np.arange(0, 12).reshape(4, 3), chunks=(2, 3))
    # Reshaping to arr.reshape(-1, ) will return
    # [ 0  1  2  3  4  5  6  7  8  9 10 11]
    # The first dimension of the reshaped dimension are the chunks with length 2
    # Assume we split the second dimension into (2, 1), i.e. setting the chunks to
    # ((2, 2), (2, 1)) and the output chunks to ((4, 2, 4, 2), )
    # In this case, the individual chunks do not hold a contiguous sequence.
    # For example, the first chunk is [[0, 1], [3, 4]].
    # Then, the result will be different because we first reshape the individual, non-contiguous chunks before concatenating them:
    # [ 0  1  3  4  2  5  6  7  9 10  8 11]
    # This is equivalent to
    # arr = np.arange(0, 12).reshape(4, 3)
    # np.concatenate(list(map(lambda x: x.reshape(-1), [arr[:2, :2], arr[:2, 2:], arr[2:, :2], arr[2:, 2:]])))

    ileft_orig = ileft
    max_result_in_chunk = _cal_max_chunk_size(result_inchunks, ileft, ii)
    if max_in_chunk == max_result_in_chunk:
        # reshaping doesn't mess up
        return result_inchunks

    while all(x == 1 for x in result_inchunks[ileft]):
        # Find the first dimension where we can split chunks
        ileft += 1

    if ileft < ii + 1:
        factor = math.ceil(max_result_in_chunk / max_in_chunk)
        result_in_chunk = result_inchunks[ileft]

        if len(result_in_chunk) == 1:

Review comment (suggestion):
```suggestion
    # The first dimension of the reshaped axis are the chunks with length 2
```

Code after:
def _calc_lower_dimension_chunks(chunks, start, stop):
    # We need the lower dimension chunks to match what the higher dimension chunks
    # can be combined to, i.e. multiply the different dimensions
    return tuple(
        map(
            lambda x: reduce(mul, x),
            product(*chunks[start : stop + 1]),
        )
    )


def _smooth_chunks(ileft, ii, max_in_chunk, result_inchunks):
    # The previous step squashed the whole dimension into a single
    # chunk for ileft + 1 (and potentially combined too many elements
    # into a single chunk for ileft as well). We split up the single
    # chunk into multiple chunks to match the max_in_chunk to keep
    # chunksizes consistent:
    # ((1, 1), (200)) -> ((1, 1), (20, ) * 10) for max_in_chunk = 20
    # It's important to ensure that all dimensions before the dimension
    # we adjust have all-1 chunks to respect C contiguous arrays
    # during the reshaping
    # Example:
    # Assume arr = da.from_array(np.arange(0, 12).reshape(4, 3), chunks=(2, 3))
    # Reshaping to arr.reshape(-1, ) will return
    # [ 0  1  2  3  4  5  6  7  8  9 10 11]
    # The first dimension of the reshaped axis are the chunks with length 2
    # Assume we split the second dimension into (2, 1), i.e. setting the chunks to
    # ((2, 2), (2, 1)) and the output chunks to ((4, 2, 4, 2), )
    # In this case, the individual chunks do not hold a contiguous sequence.
    # For example, the first chunk is [[0, 1], [3, 4]].
    # Then, the result will be different because we first reshape the individual,
    # non-contiguous chunks before concatenating them:
    # [ 0  1  3  4  2  5  6  7  9 10  8 11]
    # This is equivalent to
    # arr = np.arange(0, 12).reshape(4, 3)
    # np.concatenate(list(map(lambda x: x.reshape(-1), [arr[:2, :2], arr[:2, 2:], arr[2:, :2], arr[2:, 2:]])))

    ileft_orig = ileft
    max_result_in_chunk = _cal_max_chunk_size(result_inchunks, ileft, ii)
    if max_in_chunk == max_result_in_chunk:
        # reshaping doesn't mess up
        return result_inchunks

    while all(x == 1 for x in result_inchunks[ileft]):
        # Find the first dimension where we can split chunks
        ileft += 1

    if ileft < ii + 1:
        factor = math.ceil(max_result_in_chunk / max_in_chunk)
        result_in_chunk = result_inchunks[ileft]