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

Code before:

    Parameters
    ----------
    x: Array
        The input array to reshape.
    shape : int or tuple of ints
        The new shape should be compatible with the original shape. If
        an integer, then the result will be a 1-D array of that length.
        One shape dimension can be -1. In this case, the value is
        inferred from the length of the array and remaining dimensions.
    chunks: tuple of ints, default None
        The chunk sizes for every chunk in the output array. Dask will expand
        the chunks per dimension into the cross product of chunks for every
        chunk in the array.

        .. note::
            This information is required if the number of dimensions is increased.
            Dask cannot infer the output chunks in this case. The keyword is ignored
            if the number of dimensions is reduced.

    Notes
    -----
    This is a parallelized version of the ``np.reshape`` function with the
    following limitations:

    1.  It does not return elements in the same order as NumPy would
    2.  It only allows for reshapings that collapse like
        ``(1, 2, 3, 4) -> (1, 6, 4)``

    Examples
    --------
    >>> import dask.array as da
    >>> import numpy as np
    >>> x = da.from_array(np.arange(0, 27).reshape(3, 3, 3), chunks=(3, 2, (2, 1)))
    >>> result = reshape_blockwise(x, (3, 9))
    >>> result.chunks
    ((3,), (4, 2, 2, 1))

    The resulting chunks are calculated automatically to match the new shape.

    >>> result.compute()
    array([[ 0,  1,  3,  4,  2,  5,  6,  7,  8],
           [ 9, 10, 12, 13, 11, 14, 15, 16, 17],
           [18, 19, 21, 22, 20, 23, 24, 25, 26]])

    >>> result = reshape_blockwise(x, (3, 3, 3), chunks=x.chunks)
    >>> result.chunks
    ((3,), (2, 1), (2, 1))

    The resulting chunks are taken from the input. Chaining the reshape operation
    together like this reverts the previous reshaping operation that reduces the

Review comment (suggestion):
From a UX perspective, I think it might make more sense to raise (or warn) if `chunks` is given but the number of dimensions is reduced. Otherwise, this might lead to some unexpected behavior where `chunks` are passed but the output chunks don't match that.

Code after:
    Blockwise-reshape reshapes every block into the new shape and concatenates
    the results. This is a trivial blockwise computation but will return the
    result in a different order than NumPy. This is a good solution for
    subsequent operations that don't rely on the order.

    Parameters
    ----------
    x: Array
        The input array to reshape.
    shape : int or tuple of ints
        The new shape should be compatible with the original shape. If
        an integer, then the result will be a 1-D array of that length.
        One shape dimension can be -1. In this case, the value is
        inferred from the length of the array and remaining dimensions.
    chunks: tuple of ints, default None
        The chunk sizes for every chunk in the output array. Dask will expand
        the chunks per dimension into the cross product of chunks for every
        chunk in the array.

        An error is raised if chunks is given and the number of dimensions
        decreases.

        .. note::
            This information is required if the number of dimensions is increased.
            Dask cannot infer the output chunks in this case. The keyword is ignored
            if the number of dimensions is reduced.

    Notes
    -----
    This is a parallelized version of the ``np.reshape`` function with the
    following limitations:

    1.  It does not return elements in the same order as NumPy would
    2.  It only allows for reshapings that collapse like
        ``(1, 2, 3, 4) -> (1, 6, 4)``

    Examples
    --------
    >>> import dask.array as da
    >>> import numpy as np
    >>> x = da.from_array(np.arange(0, 27).reshape(3, 3, 3), chunks=(3, 2, (2, 1)))
    >>> result = reshape_blockwise(x, (3, 9))
    >>> result.chunks
    ((3,), (4, 2, 2, 1))

    The resulting chunks are calculated automatically to match the new shape.

    >>> result.compute()
    array([[ 0,  1,  3,  4,  2,  5,  6,  7,  8],
           [ 9, 10, 12, 13, 11, 14, 15, 16, 17],
           [18, 19, 21, 22, 20, 23, 24, 25, 26]])