dask.array.reshape_blockwise

dask.array.reshape_blockwise

dask.array.reshape_blockwise(x: dask.array.core.Array, shape: int | tuple[int, ...], chunks: tuple[tuple[int, ...], ...] | None = None) dask.array.core.Array[source]

Blockwise-reshape into a new shape.

The regular reshape operation in Dask preserves C-ordering in the array which requires a rechunking for most reshaping operations, making the computation relatively expensive.

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.

shapeint 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]])
>>> result = reshape_blockwise(result, (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 number of dimensions.

>>> result.compute()
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8]],

       [[ 9, 10, 11],
        [12, 13, 14],
        [15, 16, 17]],

       [[18, 19, 20],
        [21, 22, 23],
        [24, 25, 26]]])