Repository: dask/dask
Pull request #11337 (2024-08-21 08:15:45)

Choose automatically between tasks-based and p2p rechunking
This PR adds a heuristic to pick the "best" rechunking method between `tasks` and `p2p`.

The goal of this heuristic is to ensure for rechunking to be stable, not to pick the fastest method.

- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

     13 files  ±0       13 suites  ±0   3h 1m 19s :stopwatch: - 1m 47s
 13 164 tests ±0   12 104 :white_check_mark:  -  4   1 056 :zzz: ±0   4 :x: + 4 
137 173 runs  ±0  118 010 :white_check_mark:  - 44  19 119 :zzz: ±0  44 :x: +44 

For more details on these failures, see [this check](https://github.com/dask/dask/runs/29045776131).

Results for commit a21d7493. ± Comparison against base commit cda582c2.

[test-results]:data:application/gzip;base64,H4sIAG+sxWYC/02MSQ6DMAwAv4Jy7sEOkKWfqUJIpKgsVUhOqH+vKQSQfJmxPSvzYXALe1ZYPyq25JAu6nM0KczTxqCkJkXbtO9RNIVfS7Z2kxzhJt/h8/9sxem8CQO568jFOEcyQCbmaU9LlPXBZxoVIJzySGtEXVxJN0Xc03Yex5AImOHYy0bX3gJXPYDjIIW1NKrrlGzBGIeeC/b9AdxkXmMbAQAA

hendrikmakait: Benchmarking shows that this PR significantly reduces the runtime for all-to-all rechunking and may improve performance for more localized rechunking (such as `test_tiles_to_rows`). 

<img width="1673" alt="Screenshot 2024-08-26 at 14 41 17" src="https://github.com/user-attachments/assets/87079838-ad00-4ee8-afbb-afdf90394b35">

We also see that this PR keeps the memory low and stable for both of these tests

<img width="1673" alt="Screenshot 2024-08-26 at 14 41 27" src="https://github.com/user-attachments/assets/5aa10657-084f-4103-8c2c-0418e0518e46">

Looking at the the performance of between `test_adjacent_groups[auto-tasks]` and `test_adjacent_groups[auto-p2p-disk]`, which still defaults to `tasks` after this PR, I estimate a performance degradation caused by this heuristic to have an upper limit of ~35%. This can be counteracted by manually choosing a preferred rechunking method.
phofl: cc @dcherian 

Patch:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
index 0874199cc..5ffbf73f9 100644
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -135,7 +135,7 @@ try:
         trace,
         var,
     )
-    from dask.array.reshape import reshape
+    from dask.array.reshape import reshape, reshape_blockwise
     from dask.array.routines import (
         allclose,
         append,
diff --git a/dask/array/rechunk.py b/dask/array/rechunk.py
index 7f45fabcf..637185ecf 100644
--- a/dask/array/rechunk.py
+++ b/dask/array/rechunk.py
@@ -361,7 +361,8 @@ def rechunk(
 
     _validate_rechunk(x.chunks, chunks)
 
-    method = method or config.get("array.rechunk.method")
+    if method is None:
+        method = _choose_rechunk_method(x.chunks, chunks, threshold=threshold)
 
     if method == "tasks":
         steps = plan_rechunk(
@@ -387,6 +388,23 @@ def rechunk(
         raise NotImplementedError(f"Unknown rechunking method '{method}'")
 
 
+def _choose_rechunk_method(old_chunks, new_chunks, threshold=None):
+    if method := config.get("array.rechunk.method", None):
+        return method
+    try:
+        from distributed import default_client
+
+        default_client()
+    except (ImportError, ValueError):
+        return "tasks"
+
+    _old_to_new = old_to_new(old_chunks, new_chunks)
+    graph_size = math.prod(sum(len(ins) for ins in axis) for axis in _old_to_new)
+    threshold = threshold or config.get("array.rechunk.threshold")
+    graph_size_threshold = _graph_size_threshold(old_chunks, new_chunks, threshold)
+    return "tasks" if graph_size < graph_size_threshold else "p2p"
+
+
 def _number_of_blocks(chunks):
     return reduce(mul, map(len, chunks))
 
@@ -621,9 +639,7 @@ def plan_rechunk(
     block_size_limit = max([block_size_limit, largest_old_block, largest_new_block])
 
     # The graph size above which to optimize
-    graph_size_threshold = threshold * (
-        _number_of_blocks(old_chunks) + _number_of_blocks(new_chunks)
-    )
+    graph_size_threshold = _graph_size_threshold(old_chunks, new_chunks, threshold)
 
     current_chunks = old_chunks
     first_pass = True
@@ -660,6 +676,10 @@ def plan_rechunk(
     return steps + [new_chunks]
 
 
+def _graph_size_threshold(old_chunks, new_chunks, threshold):
+    return threshold * (_number_of_blocks(old_chunks) + _number_of_blocks(new_chunks))
+
+
 def _compute_rechunk(x, chunks):
     """Compute the rechunk of *x* to the given *chunks*."""
     if x.size == 0:
diff --git a/dask/array/reshape.py b/dask/array/reshape.py
index 56e1f3383..1cc79090a 100644
--- a/dask/array/reshape.py
+++ b/dask/array/reshape.py
@@ -29,17 +29,19 @@ directly.
 """
 
 
-def reshape_rechunk(inshape, outshape, inchunks):
+def reshape_rechunk(inshape, outshape, inchunks, disallow_dimension_expansion=False):
     assert all(isinstance(c, tuple) for c in inchunks)
     ii = len(inshape) - 1
     oi = len(outshape) - 1
     result_inchunks = [None for i in range(len(inshape))]
     result_outchunks = [None for i in range(len(outshape))]
+    mapper_in, one_dimensions = {}, []
 
     while ii >= 0 or oi >= 0:
         if inshape[ii] == outshape[oi]:
             result_inchunks[ii] = inchunks[ii]
             result_outchunks[oi] = inchunks[ii]
+            mapper_in[ii] = oi
             ii -= 1
             oi -= 1
             continue
@@ -50,13 +52,18 @@ def reshape_rechunk(inshape, outshape, inchunks):
             ii -= 1
         elif dout == 1:
             result_outchunks[oi] = (1,)
+            one_dimensions.append(oi)
             oi -= 1
         elif din < dout:  # (4, 4, 4) -> (64,)
             ileft = ii - 1
+            mapper_in[ii] = oi
             while (
                 ileft >= 0 and reduce(mul, inshape[ileft : ii + 1]) < dout
             ):  # 4 < 64, 4*4 < 64, 4*4*4 == 64
+                mapper_in[ileft] = oi
                 ileft -= 1
+
+            mapper_in[ileft] = oi
             if reduce(mul, inshape[ileft : ii + 1]) != dout:
                 raise NotImplementedError(_not_implemented_message)
             # Special case to avoid intermediate rechunking:
@@ -87,6 +94,10 @@ def reshape_rechunk(inshape, outshape, inchunks):
             oi -= 1
             ii = ileft - 1
         elif din > dout:  # (64,) -> (4, 4, 4)
+            if disallow_dimension_expansion:
+                raise NotImplementedError(
+                    "reshape_blockwise not implemented for expanding dimensions without passing chunk hints."
+                )
             oleft = oi - 1
             while oleft >= 0 and reduce(mul, outshape[oleft : oi + 1]) < din:
                 oleft -= 1
@@ -111,7 +122,7 @@ def reshape_rechunk(inshape, outshape, inchunks):
             oi = oleft - 1
             ii -= 1
 
-    return tuple(result_inchunks), tuple(result_outchunks)
+    return tuple(result_inchunks), tuple(result_outchunks), mapper_in, one_dimensions
 
 
 def _calc_lower_dimension_chunks(chunks, start, stop):
@@ -324,14 +335,7 @@ def reshape(x, shape, merge_chunks=True, limit=None):
         missing_size = sanitize_index(x.size / reduce(mul, known_sizes, 1))
         shape = tuple(missing_size if s == -1 else s for s in shape)
 
-    if np.isnan(sum(x.shape)):
-        raise ValueError(
-            "Array chunk size or shape is unknown. shape: %s\n\n"
-            "Possible solution with x.compute_chunk_sizes()" % str(x.shape)
-        )
-
-    if reduce(mul, shape, 1) != x.size:
-        raise ValueError("total size of new array must be unchanged")
+    _sanity_checks(x, shape)
 
     if x.shape == shape:
         return x
@@ -353,7 +357,7 @@ def reshape(x, shape, merge_chunks=True, limit=None):
     if not merge_chunks and din > dout:
         x = x.rechunk({i: 1 for i in range(din - dout)})
 
-    inchunks, outchunks = reshape_rechunk(x.shape, shape, x.chunks)
+    inchunks, outchunks, _, _ = reshape_rechunk(x.shape, shape, x.chunks)
     x2 = x.rechunk(inchunks)
 
     # Construct graph
@@ -364,3 +368,210 @@ def reshape(x, shape, merge_chunks=True, limit=None):
 
     graph = HighLevelGraph.from_collections(name, dsk, dependencies=[x2])
     return Array(graph, name, outchunks, meta=meta)
+
+
+def _sanity_checks(x, shape):
+    if np.isnan(sum(x.shape)):
+        raise ValueError(
+            "Array chunk size or shape is unknown. shape: %s\n\n"
+            "Possible solution with x.compute_chunk_sizes()" % str(x.shape)
+        )
+
+    if reduce(mul, shape, 1) != x.size:
+        raise ValueError("total size of new array must be unchanged")
+
+
+def reshape_blockwise(
+    x: Array,
+    shape: int | tuple[int, ...],
+    chunks: tuple[tuple[int, ...], ...] | None = None,
+) -> Array:
+    """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.
+    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]])
+
+    >>> 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]],
+    <BLANKLINE>
+           [[ 9, 10, 11],
+            [12, 13, 14],
+            [15, 16, 17]],
+    <BLANKLINE>
+           [[18, 19, 20],
+            [21, 22, 23],
+            [24, 25, 26]]])
+    """
+    if shape in [-1, (-1,)]:
+        shape = (reduce(mul, x.shape),)
+
+    if not isinstance(shape, tuple):
+        shape = (shape,)
+
+    _sanity_checks(x, shape)
+
+    outname = "reshape-blockwise-" + tokenize(x, shape)
+    chunk_tuples = list(product(*(range(len(c)) for i, c in enumerate(x.chunks))))
+
+    if len(shape) > x.ndim:
+        if chunks is None:
+            raise TypeError("Need to specify chunks if expanding dimensions.")
+        out_shapes = list(product(*(c for c in chunks)))
+        out_chunk_tuples = list(product(*(range(len(c)) for c in chunks)))
+        in_shapes = list(product(*(c for c in x.chunks)))
+        non_matching_chunks = [
+            (i, in_c, out_c)
+            for i, (in_c, out_c) in enumerate(zip(in_shapes, out_shapes))
+            if math.prod(in_c) != math.prod(out_c)
+        ]
+        if non_matching_chunks:
+            raise ValueError(
+                f"Chunk sizes do not match for the following chunks: "
+                f"{[x[0] for x in non_matching_chunks[:5]]}. \n"
+                f"The corresponding chunksizes are: {[x[1:] for x in non_matching_chunks[:5]]}. "
+                f"(restricted to first 5 entries)."
+            )
+
+        dsk = {
+            (outname,)
+            + tuple(chunk_out): (_reshape_blockwise, (x.name,) + tuple(chunk_in), shape)
+            for chunk_in, chunk_out, shape in zip(
+                chunk_tuples, out_chunk_tuples, out_shapes
+            )
+        }
+
+        graph = HighLevelGraph.from_collections(outname, dsk, dependencies=[x])  # type: ignore[arg-type]
+        return Array(graph, outname, chunks, meta=x._meta)
+
+    if chunks is not None:
+        raise ValueError(
+            "Setting chunks is not allowed when reducing the number of dimensions."
+        )
+
+    _, _, mapper_in, one_dimensions = reshape_rechunk(
+        x.shape, shape, x.chunks, disallow_dimension_expansion=True
+    )
+
+    if len(shape) == x.ndim:
+        return Array(x.dask, x.name, x.chunks, meta=x)
+
+    # Convert input chunks to output chunks
+    out_shapes = [
+        _convert_to_shape(c, mapper_in, one_dimensions)
+        for c in list(product(*(c for c in x.chunks)))
+    ]
+    # Calculate the number of chunks for each dimension in the output
+    nr_out_chunks = _convert_to_shape(
+        tuple(map(len, x.chunks)), mapper_in, one_dimensions
+    )
+    # Create output chunk tuples for the graph
+    out_chunk_tuples = list(product(*(range(c) for c in nr_out_chunks)))
+
+    dsk = {
+        (outname,)
+        + tuple(chunk_out): (_reshape_blockwise, (x.name,) + tuple(chunk_in), shape)
+        for chunk_in, chunk_out, shape in zip(
+            chunk_tuples, out_chunk_tuples, out_shapes
+        )
+    }
+
+    # Calculating the output chunks is a bit tricky. We have all the output chunk
+    # tuples, but now we have to create the source of the cross product. We can
+    # iterate over the tuples to extract all elements along a single dimension.
+    output_chunks = []
+    ctr = 1
+    for i, nr_chunks_dim in enumerate(reversed(nr_out_chunks)):
+        dimension_chunks = []
+        for elem in range(nr_chunks_dim):
+            dimension_chunks.append(out_shapes[elem * ctr][len(nr_out_chunks) - i - 1])
+        output_chunks.append(tuple(dimension_chunks))
+        ctr *= nr_chunks_dim
+
+    graph = HighLevelGraph.from_collections(outname, dsk, dependencies=[x])  # type: ignore[arg-type]
+    return Array(graph, outname, tuple(reversed(output_chunks)), meta=x._meta)
+
+
+def _reshape_blockwise(arr, shape):
+    return arr.reshape(shape)
+
+
+def _convert_to_shape(
+    shape: tuple[int, ...], mapper_in: dict[int, int], one_dims: list[int]
+):
+    # Used to map the input dimensions onto output dimensions. The mapper tracks
+    # which input dimensions are mapped onto what output dimensions.
+    # i.e. {0: 0, 1:0, 2: 1} means that dimensions 0 and 1 are mapped into 0 and 2 into 1
+    output_shape: list[list[int]] = [[]] * (
+        len(set(mapper_in.values())) + len(one_dims)
+    )
+    output_shape = list(map(lambda x: x.copy(), output_shape))
+    for i in one_dims:
+        output_shape[i] = [1]
+
+    for k, v in mapper_in.items():
+        output_shape[v].append(shape[k])
+
+    return tuple(map(lambda x: reduce(mul, x), output_shape))
diff --git a/dask/dask-schema.yaml b/dask/dask-schema.yaml
index 548c8e964..4be821187 100644
--- a/dask/dask-schema.yaml
+++ b/dask/dask-schema.yaml
@@ -131,10 +131,12 @@ properties:
         properties:
           
           method:
-            type: string
+            type:
+            - string
+            - "null"
             description: |
-              The method to use for rechunking. Must be either "tasks" or "p2p";
-              default is "tasks". Using "p2p" requires a distributed cluster.
+              The method to use for rechunking. Possible values are tasks or p2p.
+              If null, pick best method depending on application.
           
           threshold:
             type: integer
diff --git a/dask/dask.yaml b/dask/dask.yaml
index e8d00aa14..5ba27ce89 100644
--- a/dask/dask.yaml
+++ b/dask/dask.yaml
@@ -23,7 +23,7 @@ array:
   chunk-size: "128MiB"
   chunk-size-tolerance: 1.25  # Tolerance for different algorithms when creating output chunks.
   rechunk:
-    method: "tasks"  # Rechunking method to use
+    method: null  # Rechunking method to use
     threshold: 4
   svg:
     size: 120  # pixels
diff --git a/docs/source/array-api.rst b/docs/source/array-api.rst
index 4765378e5..7c6f4897e 100644
--- a/docs/source/array-api.rst
+++ b/docs/source/array-api.rst
@@ -197,6 +197,7 @@ Top level functions
    remainder
    repeat
    reshape
+   reshape_blockwise
    result_type
    right_shift
    rint


Test patch:
diff --git a/dask/array/tests/test_reshape.py b/dask/array/tests/test_reshape.py
index bc9dfa461..14fb84ba9 100644
--- a/dask/array/tests/test_reshape.py
+++ b/dask/array/tests/test_reshape.py
@@ -4,10 +4,12 @@ import numpy as np
 import pytest
 
 import dask.array as da
+from dask.array import from_array
 from dask.array.reshape import (
     _smooth_chunks,
     contract_tuple,
     expand_tuple,
+    reshape_blockwise,
     reshape_rechunk,
 )
 from dask.array.utils import assert_eq
@@ -53,7 +55,7 @@ from dask.array.utils import assert_eq
     ],
 )
 def test_reshape_rechunk(inshape, outshape, prechunks, inchunks, outchunks):
-    result_in, result_out = reshape_rechunk(inshape, outshape, prechunks)
+    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)))
@@ -127,7 +129,7 @@ def test_reshape_all_chunked_no_merge(inshape, inchunks, outshape, outchunks):
     a = da.from_array(base, chunks=inchunks)
 
     # test directly
-    inchunks2, outchunks2 = reshape_rechunk(a.shape, outshape, inchunks)
+    inchunks2, outchunks2, _, _ = reshape_rechunk(a.shape, outshape, inchunks)
     assert inchunks2 == inchunks
     assert outchunks2 == outchunks
 
@@ -165,7 +167,7 @@ def test_reshape_all_not_chunked_merge(
     a = da.from_array(base, chunks=inchunks)
 
     # test directly
-    inchunks2, outchunks2 = reshape_rechunk(a.shape, outshape, inchunks)
+    inchunks2, outchunks2, _, _ = reshape_rechunk(a.shape, outshape, inchunks)
     assert inchunks2 == expected_inchunks
     assert outchunks2 == outchunks
 
@@ -325,8 +327,120 @@ def test_reshape_split_out_chunks():
     )
 
 
+def test_reshape_blockwise():
+    x = np.arange(0, 54).reshape(6, 3, 3)
+    arr = from_array(x, chunks=(3, 2, (2, 1)))
+    assert arr.chunks == ((3, 3), (2, 1), (2, 1))
+    result = reshape_blockwise(arr, (18, 3))
+    assert result.chunks == ((6, 3, 6, 3), (2, 1))
+    expected = [
+        np.concatenate(
+            [x[:3, :2, :2].reshape(6, 2), x[:3, :2, 2:].reshape(6, 1)], axis=1
+        ),
+        np.concatenate(
+            [x[:3, 2:, :2].reshape(3, 2), x[:3, 2:, 2:].reshape(3, 1)], axis=1
+        ),
+        np.concatenate(
+            [x[3:, :2, :2].reshape(6, 2), x[3:, :2, 2:].reshape(6, 1)], axis=1
+        ),
+        np.concatenate(
+            [x[3:, 2:, :2].reshape(3, 2), x[3:, 2:, 2:].reshape(3, 1)], axis=1
+        ),
+    ]
+    expected = np.concatenate(expected, axis=0)
+    assert_eq(result, expected)
+    result2 = reshape_blockwise(result, (6, 3, 3), chunks=arr.chunks)
+    assert_eq(arr, result2)
+    assert result2.chunks == arr.chunks
+
+    result = reshape_blockwise(arr, (6, 9))
+    assert result.chunks == ((3, 3), (4, 2, 2, 1))
+    expected = [
+        np.concatenate(
+            [
+                x[:3, :2, :2].reshape(3, 4),
+                x[:3, :2, 2:].reshape(3, 2),
+                x[:3, 2:, :2].reshape(3, 2),
+                x[:3, 2:, 2:].reshape(3, 1),
+            ],
+            axis=1,
+        ),
+        np.concatenate(
+            [
+                x[3:, :2, :2].reshape(3, 4),
+                x[3:, :2, 2:].reshape(3, 2),
+                x[3:, 2:, :2].reshape(3, 2),
+                x[3:, 2:, 2:].reshape(3, 1),
+            ],
+            axis=1,
+        ),
+    ]
+
+    # chunks of first half of the array, i.e. arr[:3, :, :]:
+    # [[[0, 1], [3, 4]], [[9, 10], [12, 13]], [[18, 19], [21, 22]]]
+    # [[[2], [5]], [[11], [14]], [[20], [23]]]
+    # [[[6, 7], [15, 16]], [[24, 25]]]
+    # [[[8], [17]], [[26]]]
+    # The inner dimensions are squashed, i.e.
+    # [[ 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]]
+    expected = np.concatenate(expected, axis=0)
+    assert_eq(result, expected)
+    result2 = reshape_blockwise(result, (6, 3, 3), chunks=arr.chunks)
+    assert_eq(arr, result2)
+    assert result2.chunks == arr.chunks
+
+    with pytest.raises(
+        ValueError, match="Chunk sizes do not match for the following chunks"
+    ):
+        reshape_blockwise(result, (6, 3, 3), chunks=((3, 4), (3, 1), (3, 1)))
+
+
+def test_reshape_blockwise_raises_for_expansion():
+    x = np.arange(0, 54).reshape(6, 3, 3)
+    arr = from_array(x, chunks=(3, 2, (2, 1)))
+    msg = "Need to specify chunks if expanding dimensions."
+    with pytest.raises(TypeError, match=msg):
+        reshape_blockwise(arr, (2, 3, 3, 3))
+
+    with pytest.raises(
+        NotImplementedError,
+        match="reshape_blockwise not implemented for expanding dimensions",
+    ):
+        reshape_blockwise(arr, (3, 2, 9))
+
+
+@pytest.mark.parametrize("shape", [(18, 3), (6, 3, 3)])
+def test_reshape_blockwise_dimension_reduction_and_chunks(shape):
+    x = np.arange(0, 54).reshape(6, 3, 3)
+    arr = from_array(x, chunks=(3, 2, (2, 1)))
+    with pytest.raises(ValueError, match="Setting chunks is not allowed when"):
+        reshape_blockwise(arr, shape, arr.chunks)
+
+
+def test_reshape_blockwise_identity():
+    x = np.arange(0, 54).reshape(6, 3, 3)
+    arr = from_array(x, chunks=(3, 2, (2, 1)))
+    result = reshape_blockwise(arr, (6, 3, 3))
+    assert len(arr.dask) == len(result.dask)
+    assert_eq(result, x)
+
+
 def test_argwhere_reshaping_not_concating_chunks():
     # GH#10080
     arr = da.random.random((500, 500, 500), chunks=(100, 100, 100)) < 0
     result = da.argwhere(arr)
     assert result.chunks == ((np.nan,) * 125, (1, 1, 1))
+
+
+def test_blockwise_reshape_sanity_checks():
+    x = np.arange(0, 20).reshape(10, 2)
+    arr = from_array(x, chunks=(5, 1))
+    with pytest.raises(ValueError, match="total size of new array must be unchanged"):
+        reshape_blockwise(arr, (10, 6))
+
+    arr = from_array(x, chunks=(5, 1))
+    arr = arr[arr[:, 0] > 0]
+    with pytest.raises(ValueError, match="Array chunk size or shape is unknown"):
+        reshape_blockwise(arr, (5, 2, 2))