Repository: dask/dask
Pull request #11542 (2024-11-20 09:12:42)

Use BlockwiseDep to implement map_blocks keywords
- [ ] Closes #xxxx
- [ ] 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 23m 21s ⏱️ + 1m 33s
 13 297 tests ±0   12 235 ✅ ±0   1 062 💤 ±0  0 ❌ ±0 
138 778 runs  ±0  119 785 ✅ ±0  18 993 💤 ±0  0 ❌ ±0 

Results for commit 77d22dca. ± Comparison against base commit 5b115c43.

[test-results]:data:application/gzip;base64,H4sIAL6xPWcC/02MzQ6DIBAGX6Xh3AO7iEBfpsEFE1J/GoRT03cvVkUPe5j5NvNhfRj8wh43EPcbW3JIJ7kcbQrztDIih6LKmrYdjTr4uWSi7UnIU77Ce5W8xep6G4bieBU+xjnuJuZpS2ul9M41DUZpWeWe1saIw13Sf76WaR7HkAowpRyiIwsOiSTvyjWi860AT1y2HXfKNwCWfX8yFxiFGgEAAA==


Patch:
diff --git a/dask/array/blockwise.py b/dask/array/blockwise.py
index 300120e8b..0bfe4412b 100644
--- a/dask/array/blockwise.py
+++ b/dask/array/blockwise.py
@@ -9,6 +9,7 @@ from dask import base, utils
 from dask.blockwise import blockwise as core_blockwise
 from dask.delayed import unpack_collections
 from dask.highlevelgraph import HighLevelGraph
+from dask.layers import ArrayBlockwiseDep
 
 
 def blockwise(
@@ -224,9 +225,10 @@ def blockwise(
                     "Index string %s does not match array dimension %d"
                     % (ind, arg.ndim)
                 )
-            numblocks[arg.name] = arg.numblocks
-            arrays.append(arg)
-            arg = arg.name
+            if not isinstance(arg, ArrayBlockwiseDep):
+                numblocks[arg.name] = arg.numblocks
+                arrays.append(arg)
+                arg = arg.name
         argindsstr.extend((arg, ind))
 
     # Normalize keyword arguments
diff --git a/dask/array/core.py b/dask/array/core.py
index 10736af6f..6a255c76d 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -56,7 +56,7 @@ from dask.context import globalmethod
 from dask.core import quote
 from dask.delayed import Delayed, delayed
 from dask.highlevelgraph import HighLevelGraph, MaterializedLayer
-from dask.layers import ArraySliceDep, reshapelist
+from dask.layers import ArrayBlockIdDep, ArraySliceDep, ArrayValuesDep, reshapelist
 from dask.sizeof import sizeof
 from dask.typing import Graph, Key, NestedKeys
 from dask.utils import (
@@ -887,57 +887,25 @@ def map_blocks(
 
     extra_argpairs = []
     extra_names = []
-    # If func has block_id as an argument, construct an array of block IDs and
-    # prepare to inject it.
-
-    def _getter_item(a, b, **kwargs):
-        if len(a.shape) == 0:
-            # meta
-            return a.item()
-        return getter(a, b, **kwargs).item()
-
-    def _add_blockwise_layer_for_keyword(keyword_arr, prefix):
-        name = prefix + out.name
-        cs = tuple((1,) * len(c) for c in out.chunks)
-
-        dsk = graph_from_arraylike(
-            keyword_arr,
-            cs,
-            keyword_arr.shape,
-            name,
-            getitem=_getter_item,
-            dtype=keyword_arr.dtype,
-            inline_array=False,
-        )
-        return Array(dsk, name, chunks=cs, dtype=keyword_arr.dtype)
 
+    # If func has block_id as an argument, construct an object to inject it.
     if has_keyword(func, "block_id"):
-        # put block_id into a Blockwise layer so that we can fuse it
-        # with the other blockwise layers
-        block_id_arr = np.empty(tuple([len(c) for c in out.chunks]), dtype=np.object_)
-        for block_id in product(*(range(len(c)) for c in out.chunks)):
-            block_id_arr[block_id] = block_id
-        block_id_array = _add_blockwise_layer_for_keyword(block_id_arr, "block-id-")
-        extra_argpairs.append((block_id_array, out_ind))
+        extra_argpairs.append((ArrayBlockIdDep(out.chunks), out_ind))
         extra_names.append("block_id")
 
     if has_keyword(func, "_overlap_trim_info"):
         # Internal for map overlap to reduce size of graph
         num_chunks = out.numblocks
-        block_id_arr = np.empty(tuple([len(c) for c in out.chunks]), dtype=np.object_)
-        for block_id in product(*(range(len(c)) for c in out.chunks)):
-            block_id_arr[block_id] = (block_id, num_chunks)
-        block_id_array = _add_blockwise_layer_for_keyword(
-            block_id_arr, "_overlap_trim_info-id-"
-        )
-        extra_argpairs.append((block_id_array, out_ind))
+        block_id_dict = {
+            block_id: (block_id, num_chunks)
+            for block_id in product(*(range(len(c)) for c in out.chunks))
+        }
+        extra_argpairs.append((ArrayValuesDep(out.chunks, block_id_dict), out_ind))
         extra_names.append("_overlap_trim_info")
 
-    # If func has block_info as an argument, construct an array of block info
+    # If func has block_info as an argument, construct a dict of block info
     # objects and prepare to inject it.
     if has_keyword(func, "block_info"):
-        # put block_info into a Blockwise layer so that we can fuse it
-        # with the other blockwise layers
         starts = {}
         num_chunks = {}
         shapes = {}
@@ -964,8 +932,7 @@ def map_blocks(
                     num_chunks[i] = arg.numblocks
         out_starts = [cached_cumsum(c, initial_zero=True) for c in out.chunks]
 
-        block_info_name = "block-info-" + out.name
-        block_info_arr = np.empty(tuple([len(c) for c in out.chunks]), dtype=np.object_)
+        block_info_dict = {}
         for block_id in product(*(range(len(c)) for c in out.chunks)):
             # Get position of chunk, indexed by axis labels
             location = {out_ind[i]: loc for i, loc in enumerate(block_id)}
@@ -1002,20 +969,9 @@ def map_blocks(
                 ),
                 "dtype": dtype,
             }
-            block_info_arr[block_id] = info
+            block_info_dict[block_id] = info
 
-        cs = tuple((1,) * len(c) for c in out.chunks)
-        dsk = graph_from_arraylike(
-            block_info_arr,
-            cs,
-            block_info_arr.shape,
-            block_info_name,
-            getitem=_getter_item,
-            dtype=block_info_arr.dtype,
-            inline_array=False,
-        )
-        block_info = Array(dsk, block_info_name, chunks=cs, dtype=np.object_)
-        extra_argpairs.append((block_info, out_ind))
+        extra_argpairs.append((ArrayValuesDep(out.chunks, block_info_dict), out_ind))
         extra_names.append("block_info")
 
     if extra_argpairs:
@@ -3754,9 +3710,9 @@ def from_zarr(
                 store = zarr.storage.FSStore(url, **storage_options)
         else:
             store = url
-        z = zarr.open_array(store=store, read_only=True, path=component, **kwargs)
+        z = zarr.open_array(store=store, path=component, **kwargs)
     else:
-        z = zarr.open_array(store=url, read_only=True, path=component, **kwargs)
+        z = zarr.open_array(store=url, path=component, **kwargs)
     chunks = chunks if chunks is not None else z.chunks
     if name is None:
         name = "from-zarr-" + tokenize(z, component, storage_options, chunks, **kwargs)
diff --git a/dask/array/optimization.py b/dask/array/optimization.py
index 3ec56353a..20893e30a 100644
--- a/dask/array/optimization.py
+++ b/dask/array/optimization.py
@@ -1,36 +1,21 @@
 from __future__ import annotations
 
-from collections.abc import Callable
 from itertools import zip_longest
 from numbers import Integral
-from typing import Any
-
-import numpy as np
 
 from dask import config
 from dask._task_spec import convert_legacy_graph, fuse_linear_task_spec
-from dask.array.chunk import getitem
-from dask.array.core import getter, getter_inline, getter_nofancy
 from dask.blockwise import fuse_roots, optimize_blockwise
 from dask.core import flatten
 from dask.highlevelgraph import HighLevelGraph
-from dask.optimization import SubgraphCallable
 from dask.utils import ensure_dict
 
-# All get* functions the optimizations know about
-GETTERS = (getter, getter_nofancy, getter_inline, getitem)
-# These get* functions aren't ever completely removed from the graph,
-# even if the index should be a no-op by numpy semantics. Some array-like's
-# don't completely follow semantics, making indexing always necessary.
-GETNOREMOVE = (getter, getter_nofancy)
-
 
 def optimize(dsk, keys, **kwargs):
     """Optimize dask for array computation
 
     1.  Cull tasks not necessary to evaluate keys
-    2.  Remove full slicing, e.g. x[:]
-    3.  Inline fast functions like getitem and np.transpose
+    2.  Perform linear fusion
     """
     if not isinstance(keys, (list, set)):
         keys = [keys]
@@ -56,126 +41,6 @@ def optimize(dsk, keys, **kwargs):
     return dsk
 
 
-def _is_getter_task(
-    value,
-) -> tuple[Callable, Any, Any, bool, bool | None] | None:
-    """Check if a value in a Dask graph looks like a getter.
-
-    1. Is it a tuple with the first element a known getter.
-    2. Is it a SubgraphCallable with a single element in its
-       dsk which is a known getter.
-
-    If a getter is found, it returns a tuple with (getter, array, index, asarray, lock).
-    Otherwise it returns ``None``.
-
-    TODO: the second check is a hack to allow for slice fusion between tasks produced
-    from blockwise layers and slicing operations. Once slicing operations have
-    HighLevelGraph layers which can talk to Blockwise layers this check *should* be
-    removed, and we should not have to introspect SubgraphCallables.
-    """
-    if type(value) is not tuple:
-        return None
-    first = value[0]
-    get: Callable | None = None
-    if first in GETTERS:
-        get = first
-    # We only accept SubgraphCallables with a single sub-task right now as it's
-    # not clear which task to inspect if there is more than one, or how to resolve
-    # conflicts if they occur.
-    elif isinstance(first, SubgraphCallable) and len(first.dsk) == 1:
-        v = next(iter(first.dsk.values()))
-        if type(v) is tuple and len(v) > 1 and v[0] in GETTERS:
-            get = v[0]
-    if get is None:  # Didn't find a getter
-        return None
-
-    length = len(value)
-    if length == 3:
-        # getter defaults to asarray=True, getitem is semantically False
-        return get, value[1], value[2], get is not getitem, None
-    elif length == 5:
-        return get, *value[1:]
-
-    return None
-
-
-def optimize_slices(dsk):
-    """Optimize slices
-
-    1.  Fuse repeated slices, like x[5:][2:6] -> x[7:11]
-    2.  Remove full slices, like         x[:] -> x
-
-    See also:
-        fuse_slice_dict
-    """
-    fancy_ind_types = (list, np.ndarray)
-    dsk = dsk.copy()
-    for k, v in dsk.items():
-        if a_task := _is_getter_task(v):
-            get, a, a_index, a_asarray, a_lock = a_task
-
-            while b_task := _is_getter_task(a):
-                f2, b, b_index, b_asarray, b_lock = b_task
-
-                if a_lock and a_lock is not b_lock:
-                    break
-                if (type(a_index) is tuple) != (type(b_index) is tuple):
-                    break
-                if type(a_index) is tuple:
-                    indices = b_index + a_index
-                    if len(a_index) != len(b_index) and any(i is None for i in indices):
-                        break
-                    if f2 is getter_nofancy and any(
-                        isinstance(i, fancy_ind_types) for i in indices
-                    ):
-                        break
-                elif f2 is getter_nofancy and (
-                    type(a_index) in fancy_ind_types or type(b_index) in fancy_ind_types
-                ):
-                    break
-                try:
-                    c_index = fuse_slice(b_index, a_index)
-                    # rely on fact that nested gets never decrease in
-                    # strictness e.g. `(getter_nofancy, (getter, ...))` never
-                    # happens
-                    get = getter if f2 is getter_inline else f2
-                except NotImplementedError:
-                    break
-                a, a_index, a_lock = b, c_index, b_lock
-                a_asarray |= b_asarray
-
-            # Skip the get call if not from from_array and nothing to do
-            if get not in GETNOREMOVE and (
-                (
-                    type(a_index) is slice
-                    and not a_index.start
-                    and a_index.stop is None
-                    and a_index.step is None
-                )
-                or (
-                    type(a_index) is tuple
-                    and all(
-                        type(s) is slice
-                        and not s.start
-                        and s.stop is None
-                        and s.step is None
-                        for s in a_index
-                    )
-                )
-            ):
-                dsk[k] = a
-            elif get is getitem or (a_asarray and not a_lock):
-                # default settings are fine, drop the extra parameters Since we
-                # always fallback to inner `get` functions, `get is getitem`
-                # can only occur if all gets are getitem, meaning all
-                # parameters must be getitem defaults.
-                dsk[k] = (get, a, a_index)
-            else:
-                dsk[k] = (get, a, a_index, a_asarray, a_lock)
-
-    return dsk
-
-
 def normalize_slice(s):
     """Replace Nones in slices with integers
 
diff --git a/dask/array/reshape.py b/dask/array/reshape.py
index 1cc79090a..b824110c7 100644
--- a/dask/array/reshape.py
+++ b/dask/array/reshape.py
@@ -473,6 +473,9 @@ def reshape_blockwise(
 
     _sanity_checks(x, shape)
 
+    if len(shape) == x.ndim and shape == x.shape:
+        return Array(x.dask, x.name, x.chunks, meta=x)
+
     outname = "reshape-blockwise-" + tokenize(x, shape)
     chunk_tuples = list(product(*(range(len(c)) for i, c in enumerate(x.chunks))))
 
@@ -515,9 +518,6 @@ def reshape_blockwise(
         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)
diff --git a/dask/layers.py b/dask/layers.py
index 54208579e..c0b1384f3 100644
--- a/dask/layers.py
+++ b/dask/layers.py
@@ -6,7 +6,7 @@ import operator
 from collections import defaultdict
 from collections.abc import Callable
 from itertools import product
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import tlz as toolz
 from tlz.curried import map
@@ -18,6 +18,9 @@ from dask.highlevelgraph import Layer
 from dask.tokenize import normalize_token
 from dask.utils import apply, cached_cumsum, concrete, insert
 
+if TYPE_CHECKING:
+    import numpy as np
+
 #
 ##
 ###  General Utilities
@@ -88,11 +91,38 @@ class ArraySliceDep(ArrayBlockwiseDep):
         return tuple(slice(*s, None) for s in loc)
 
 
+class ArrayBlockIdDep(ArrayBlockwiseDep):
+    def __getitem__(self, idx: tuple):
+        if not isinstance(idx, tuple):
+            # This happens if something else except graph generation calls us
+            raise NotImplementedError("ArrayBlockIdDep requires a tuple index")
+        return idx
+
+
+class ArrayValuesDep(ArrayBlockwiseDep):
+    def __init__(self, chunks: tuple[tuple[int, ...], ...], values: np.ndarray | dict):
+        super().__init__(chunks)
+        self.values = values
+
+    def __getitem__(self, idx: tuple):
+        return self.values[idx]
+
+
 @normalize_token.register(ArraySliceDep)
 def normalize_array_slice_dep(dep):
     return "ArraySliceDep", dep.chunks, dep.starts
 
 
+@normalize_token.register(ArrayBlockIdDep)
+def normalize_array_block_id_dep(dep):
+    return "ArrayBlockIdDep", dep.chunks
+
+
+@normalize_token.register(ArrayValuesDep)
+def normalize_array_values_dep(dep):
+    return "ArrayValuesDep", dep.chunks, dep.values
+
+
 class ArrayOverlapLayer(Layer):
     """Simple HighLevelGraph array overlap layer.
 


Test patch:
diff --git a/dask/array/tests/test_map_blocks.py b/dask/array/tests/test_map_blocks.py
index e88bf37eb..e27c74daf 100644
--- a/dask/array/tests/test_map_blocks.py
+++ b/dask/array/tests/test_map_blocks.py
@@ -12,4 +12,4 @@ def test_map_blocks_block_id_fusion():
 
     result = arr.map_blocks(dummy).astype("f8")
     dsk = collections_to_dsk([result])
-    assert len(dsk) == 22
+    assert len(dsk) == 20
diff --git a/dask/array/tests/test_optimization.py b/dask/array/tests/test_optimization.py
index cf8583bb5..66d25551d 100644
--- a/dask/array/tests/test_optimization.py
+++ b/dask/array/tests/test_optimization.py
@@ -13,16 +13,10 @@ import dask.array as da
 from dask.array.chunk import getitem as da_getitem
 from dask.array.core import getter as da_getter
 from dask.array.core import getter_nofancy as da_getter_nofancy
-from dask.array.optimization import (
-    fuse_slice,
-    optimize,
-    optimize_blockwise,
-    optimize_slices,
-)
+from dask.array.optimization import fuse_slice, optimize, optimize_blockwise
 from dask.array.utils import assert_eq
 from dask.highlevelgraph import HighLevelGraph
-from dask.optimization import SubgraphCallable, fuse
-from dask.utils import SerializableLock
+from dask.optimization import SubgraphCallable
 
 
 def _wrap_getter(func, wrap):
@@ -108,148 +102,6 @@ def _assert_getter_dsk_eq(a, b):
             assert av == bv
 
 
-def test_fuse_getitem(getter, getter_nofancy, getitem):
-    pairs = [
-        (
-            (getter, (getter, "x", slice(1000, 2000)), slice(15, 20)),
-            (getter, "x", slice(1015, 1020)),
-        ),
-        (
-            (
-                getitem,
-                (getter, "x", (slice(1000, 2000), slice(100, 200))),
-                (slice(15, 20), slice(50, 60)),
-            ),
-            (getter, "x", (slice(1015, 1020), slice(150, 160))),
-        ),
-        (
-            (
-                getitem,
-                (getter_nofancy, "x", (slice(1000, 2000), slice(100, 200))),
-                (slice(15, 20), slice(50, 60)),
-            ),
-            (getter_nofancy, "x", (slice(1015, 1020), slice(150, 160))),
-        ),
-        ((getter, (getter, "x", slice(1000, 2000)), 10), (getter, "x", 1010)),
-        (
-            (getitem, (getter, "x", (slice(1000, 2000), 10)), (slice(15, 20),)),
-            (getter, "x", (slice(1015, 1020), 10)),
-        ),
-        (
-            (getitem, (getter_nofancy, "x", (slice(1000, 2000), 10)), (slice(15, 20),)),
-            (getter_nofancy, "x", (slice(1015, 1020), 10)),
-        ),
-        (
-            (getter, (getter, "x", (10, slice(1000, 2000))), (slice(15, 20),)),
-            (getter, "x", (10, slice(1015, 1020))),
-        ),
-        (
-            (
-                getter,
-                (getter, "x", (slice(1000, 2000), slice(100, 200))),
-                (slice(None, None), slice(50, 60)),
-            ),
-            (getter, "x", (slice(1000, 2000), slice(150, 160))),
-        ),
-        (
-            (getter, (getter, "x", (None, slice(None, None))), (slice(None, None), 5)),
-            (getter, "x", (None, 5)),
-        ),
-        (
-            (
-                getter,
-                (getter, "x", (slice(1000, 2000), slice(10, 20))),
-                (slice(5, 10),),
-            ),
-            (getter, "x", (slice(1005, 1010), slice(10, 20))),
-        ),
-        (
-            (
-                getitem,
-                (getitem, "x", (slice(1000, 2000),)),
-                (slice(5, 10), slice(10, 20)),
-            ),
-            (getitem, "x", (slice(1005, 1010), slice(10, 20))),
-        ),
-        (
-            (getter, (getter, "x", slice(1000, 2000), False, False), slice(15, 20)),
-            (getter, "x", slice(1015, 1020)),
-        ),
-        (
-            (getter, (getter, "x", slice(1000, 2000)), slice(15, 20), False, False),
-            (getter, "x", slice(1015, 1020)),
-        ),
-        (
-            (
-                getter,
-                (getter_nofancy, "x", slice(1000, 2000), False, False),
-                slice(15, 20),
-                False,
-                False,
-            ),
-            (getter_nofancy, "x", slice(1015, 1020), False, False),
-        ),
-    ]
-
-    for inp, expected in pairs:
-        result = optimize_slices({"y": inp})
-        _assert_getter_dsk_eq(result, {"y": expected})
-
-
-def test_fuse_getitem_lock(getter, getter_nofancy, getitem):
-    lock1 = SerializableLock()
-    lock2 = SerializableLock()
-
-    pairs = [
-        (
-            (getter, (getter, "x", slice(1000, 2000), True, lock1), slice(15, 20)),
-            (getter, "x", slice(1015, 1020), True, lock1),
-        ),
-        (
-            (
-                getitem,
-                (getter, "x", (slice(1000, 2000), slice(100, 200)), True, lock1),
-                (slice(15, 20), slice(50, 60)),
-            ),
-            (getter, "x", (slice(1015, 1020), slice(150, 160)), True, lock1),
-        ),
-        (
-            (
-                getitem,
-                (
-                    getter_nofancy,
-                    "x",
-                    (slice(1000, 2000), slice(100, 200)),
-                    True,
-                    lock1,
-                ),
-                (slice(15, 20), slice(50, 60)),
-            ),
-            (getter_nofancy, "x", (slice(1015, 1020), slice(150, 160)), True, lock1),
-        ),
-        (
-            (
-                getter,
-                (getter, "x", slice(1000, 2000), True, lock1),
-                slice(15, 20),
-                True,
-                lock2,
-            ),
-            (
-                getter,
-                (getter, "x", slice(1000, 2000), True, lock1),
-                slice(15, 20),
-                True,
-                lock2,
-            ),
-        ),
-    ]
-
-    for inp, expected in pairs:
-        result = optimize_slices({"y": inp})
-        _assert_getter_dsk_eq(result, {"y": expected})
-
-
 def test_optimize_with_getitem_fusion(getter):
     dsk = {
         "a": "some-array",
@@ -262,30 +114,6 @@ def test_optimize_with_getitem_fusion(getter):
     assert len(result) < len(dsk)
 
 
-def test_optimize_slicing(getter):
-    dsk = {
-        "a": (range, 10),
-        "b": (getter, "a", (slice(None, None, None),)),
-        "c": (getter, "b", (slice(None, None, None),)),
-        "d": (getter, "c", (slice(0, 5, None),)),
-        "e": (getter, "d", (slice(None, None, None),)),
-    }
-
-    expected = {"e": (getter, (range, 10), (slice(0, 5, None),))}
-    result = optimize_slices(fuse(dsk, [], rename_keys=False)[0])
-    _assert_getter_dsk_eq(result, expected)
-
-    # protect output keys
-    expected = {
-        "c": (getter, (range, 10), (slice(0, None, None),)),
-        "d": (getter, "c", (slice(0, 5, None),)),
-        "e": (getter, "d", (slice(None, None, None),)),
-    }
-    result = optimize_slices(fuse(dsk, ["c", "d", "e"], rename_keys=False)[0])
-
-    _assert_getter_dsk_eq(result, expected)
-
-
 def test_fuse_slice():
     assert fuse_slice(slice(10, 15), slice(0, 5, 2)) == slice(10, 15, 2)
 
@@ -342,13 +170,6 @@ def test_nonfusible_fancy_indexing():
             fuse_slice(a, b)
 
 
-def test_hard_fuse_slice_cases(getter):
-    dsk = {
-        "x": (getter, (getter, "x", (None, slice(None, None))), (slice(None, None), 5))
-    }
-    _assert_getter_dsk_eq(optimize_slices(dsk), {"x": (getter, "x", (None, 5))})
-
-
 def test_dont_fuse_numpy_arrays():
     x = np.ones(10)
     for _ in [(5,), (10,)]:
@@ -376,20 +197,6 @@ def test_fuse_slices_with_alias(getter, getitem):
     assert len(dsk2) == 3
 
 
-def test_dont_fuse_fancy_indexing_in_getter_nofancy(getitem, getter_nofancy):
-    dsk = {
-        "a": (
-            getitem,
-            (getter_nofancy, "x", (slice(10, 20, None), slice(100, 200, None))),
-            ([1, 3], slice(50, 60, None)),
-        )
-    }
-    _assert_getter_dsk_eq(optimize_slices(dsk), dsk)
-
-    dsk = {"a": (getitem, (getter_nofancy, "x", [1, 2, 3]), 0)}
-    _assert_getter_dsk_eq(optimize_slices(dsk), dsk)
-
-
 @pytest.mark.parametrize("chunks", [10, 5, 3])
 def test_fuse_getter_with_asarray(chunks):
     x = np.ones(10) * 1234567890
@@ -403,50 +210,6 @@ def test_fuse_getter_with_asarray(chunks):
     assert_eq(z, x + 1)
 
 
-def test_remove_no_op_slices_for_getitem(getitem):
-    null = slice(0, None)
-    opts = [
-        ((getitem, "x", null, False, False), "x"),
-        ((getitem, (getitem, "x", null, False, False), null), "x"),
-        ((getitem, (getitem, "x", (null, null), False, False), ()), "x"),
-    ]
-    for orig, final in opts:
-        _assert_getter_dsk_eq(optimize_slices({"a": orig}), {"a": final})
-
-
-@pytest.mark.parametrize("which", ["getter", "getter_nofancy"])
-def test_dont_remove_no_op_slices_for_getter_or_getter_nofancy(
-    which, getitem, getter, getter_nofancy
-):
-    # Test that no-op slices are *not* removed when using getter or
-    # getter_nofancy. This ensures that `get` calls are always made in all
-    # tasks created by `from_array`, even after optimization
-
-    # Pytest doesn't make it easy to parameterize over parameterized fixtures
-    if which == "getter":
-        get = getter
-    else:
-        get = getter_nofancy
-
-    null = slice(0, None)
-    opts = [
-        (
-            (get, "x", null, False, False),
-            (get, "x", null, False, False),
-        ),
-        (
-            (getitem, (get, "x", null, False, False), null),
-            (get, "x", null, False, False),
-        ),
-        (
-            (getitem, (get, "x", (null, null), False, False), ()),
-            (get, "x", (null, null), False, False),
-        ),
-    ]
-    for orig, final in opts:
-        _assert_getter_dsk_eq(optimize_slices({"a": orig}), {"a": final})
-
-
 @pytest.mark.xfail(reason="blockwise fusion does not respect this, which is ok")
 def test_turn_off_fusion():
     x = da.ones(10, chunks=(5,))
diff --git a/dask/array/tests/test_overlap.py b/dask/array/tests/test_overlap.py
index 3592ccae3..4b4a35151 100644
--- a/dask/array/tests/test_overlap.py
+++ b/dask/array/tests/test_overlap.py
@@ -125,8 +125,7 @@ def test_overlap_internal_asymmetric_small():
 def test_trim_internal():
     d = da.ones((40, 60), chunks=(10, 10))
     e = trim_internal(d, axes={0: 1, 1: 2}, boundary="reflect")
-    assert any("_overlap_trim_info" in k[0] for k in dict(e.dask))
-
+    assert len(collections_to_dsk([e])) == 24
     assert e.chunks == ((8, 8, 8, 8), (6, 6, 6, 6, 6, 6))
 
 
diff --git a/dask/array/tests/test_reshape.py b/dask/array/tests/test_reshape.py
index 14fb84ba9..1473d97d1 100644
--- a/dask/array/tests/test_reshape.py
+++ b/dask/array/tests/test_reshape.py
@@ -411,12 +411,11 @@ def test_reshape_blockwise_raises_for_expansion():
         reshape_blockwise(arr, (3, 2, 9))
 
 
-@pytest.mark.parametrize("shape", [(18, 3), (6, 3, 3)])
-def test_reshape_blockwise_dimension_reduction_and_chunks(shape):
+def test_reshape_blockwise_dimension_reduction_and_chunks():
     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)
+        reshape_blockwise(arr, (18, 3), arr.chunks)
 
 
 def test_reshape_blockwise_identity():
@@ -426,6 +425,10 @@ def test_reshape_blockwise_identity():
     assert len(arr.dask) == len(result.dask)
     assert_eq(result, x)
 
+    result = reshape_blockwise(arr, (6, 3, 3), chunks=arr.chunks)
+    assert len(arr.dask) == len(result.dask)
+    assert_eq(result, x)
+
 
 def test_argwhere_reshaping_not_concating_chunks():
     # GH#10080