Repository: dask/dask
Pull request #11945 (2025-05-12 14:59:06)

Speed up slicing graph generation
This cleans up a couple of things around slicing to speed it up and restore it to its old performance. 

Closes https://github.com/dask/dask/issues/11735 (it's arguable whether it actually closes the ticket since some things are still a bit slower)

When comparing the old slicing performance with the new one, it is also important to note that the new slicing operation implements a shuffle for most nontrivial operations which is slower during graph generation but orders of magnitude better during runtime for most, if not all, affected operations.

There are many smaller performance fixes in here. However, the majority of speedup stems from disabling automatic garbage collection while this is running.


```

# Very old version
▶ python profile_slicing.py
Dask version: 2024.9.0
Empty indexer (all)                     :       0.01 ms
Every second element (::2)              :     146.57 ms
All but last element (:-1)              :     330.51 ms
Multi dim one slice (0, slice(None))    :      62.01 ms
Multi dim (0, slice(None, None, 2))     :      33.58 ms
Multi dim (0, slice(None, -1))          :      65.67 ms


# Main
(dask-distributed)
~/workspace/dask  ae250c9ac ✗                                                                                                                                                                                                         240d17h ◒
▶ python profile_slicing.py
Dask version: 2025.4.1+8.g4914f2f58
Empty indexer (all)                     :       0.01 ms
Every second element (::2)              :     467.75 ms
All but last element (:-1)              :     636.76 ms
Multi dim one slice (0, slice(None))    :     550.30 ms
Multi dim (0, slice(None, None, 2))     :     254.68 ms
Multi dim (0, slice(None, -1))          :     498.62 ms


# this PR
(dask-distributed)
~/workspace/dask  main ✗                                                                                                                                                                                                                 7d7h ◒
▶ python profile_slicing.py
Dask version: 2025.4.1+22.g2e6c80732
Empty indexer (all)                     :       0.02 ms
Every second element (::2)              :     156.41 ms
All but last element (:-1)              :     246.63 ms
Multi dim one slice (0, slice(None))    :     168.82 ms
Multi dim (0, slice(None, None, 2))     :      84.51 ms
Multi dim (0, slice(None, -1))          :     185.94 ms
```


<details>
<summary>Benchmark code</summary>

```python
import time
import dask.array as da
import dask

print(f"Dask version: {dask.__version__}")

x = da.ones((100_000,), chunks=(1,))
niterations = 5


def benchmark(arr, indexer, identifier):
    try:
        start = time.perf_counter()
        for i in range(niterations):
            arr[indexer]
        end = time.perf_counter()
        print(f"{identifier:<40}: {(end - start)/niterations * 1000:>10.2f} ms")
    except Exception as e:
        print(f"Error with {identifier}: {e}")
        return


benchmark(x, slice(None, None, None), "Empty indexer (all)")
benchmark(x, slice(None, None, 2), "Every second element (::2)")
benchmark(x, slice(None, -1), "All but last element (:-1)")

x = da.ones((100, 1000, 100), chunks=(-1, 1, 1))

benchmark(x, (1, slice(None)), "Multi dim one slice (0, slice(None))")
benchmark(x, (1, slice(None, None, 2)), "Multi dim (0, slice(None, None, 2))")
benchmark(x, (1, slice(None, -1)), "Multi dim (0, slice(None, -1))")
```

</details>

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._

      9 files  ±0        9 suites  ±0   3h 18m 27s ⏱️ - 1m 44s
 18 029 tests ±0   16 801 ✅  -  14   1 214 💤 ±0   14 ❌ + 14 
161 335 runs  ±0  149 095 ✅  - 124  12 114 💤  - 2  126 ❌ +126 

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

Results for commit 66a2810a. ± Comparison against base commit ba63b71b.

[test-results]:data:application/gzip;base64,H4sIAO0UImgC/03Myw7CIBCF4VcxrF3MwDgtvoyhFBJiL4bCyvjuQmOpy/Od5H8LHya3iftFXy9iyyG1MeZoUliXMhE1dIXKmeqNPUh97MeWra3IPeCJz/CqKJGaeROmaqe4GNdYCIrEvOxtRqVuv93apEGf2Np7arejLfmQ/7Zd5zmkMgSzkT2CYQcOlfTEAB0ZUpZAK4+esBsGNYrPF1OGpLAbAQAA

fjetter: TODO: Test if https://github.com/dask/dask/issues/11853 is fixed as well

Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 9eb0924aa..9fd4f75a7 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -354,6 +354,20 @@ class TaskRef:
             return False
         return self.key == value.key
 
+    def __reduce__(self):
+        return TaskRef, (self.key,)
+
+    def substitute(self, subs: dict, key: KeyType | None = None) -> TaskRef | GraphNode:
+        if self.key in subs:
+            val = subs[self.key]
+            if isinstance(val, GraphNode):
+                return val.substitute({}, key=self.key)
+            elif isinstance(val, TaskRef):
+                return val
+            else:
+                return TaskRef(val)
+        return self
+
 
 class GraphNode:
     key: KeyType
@@ -522,7 +536,7 @@ class Alias(GraphNode):
             val = subs.get(self.target, self.target)
             if sub_key == self.key and val == self.target:
                 return self
-            if isinstance(val, GraphNode):
+            if isinstance(val, (GraphNode, TaskRef)):
                 return val.substitute({}, key=key)
             if key is None and isinstance(sub_key, GraphNode):
                 raise RuntimeError(
@@ -599,18 +613,10 @@ class DataNode(GraphNode):
         return iter(self.value)
 
 
-def _call_recursively(obj, values):
-    if isinstance(obj, GraphNode):
-        return obj(values)
-    elif isinstance(obj, dict):
-        return {k: _call_recursively(v, values) for k, v in obj.items()}
-    elif isinstance(obj, (list, tuple, frozenset, set)):
-        return type(obj)(_call_recursively(v, values) for v in obj)
-    return obj
-
-
 def _get_dependencies(obj: object) -> set | frozenset:
-    if isinstance(obj, GraphNode):
+    if isinstance(obj, TaskRef):
+        return {obj.key}
+    elif isinstance(obj, GraphNode):
         return obj.dependencies
     elif isinstance(obj, dict):
         if not obj:
@@ -647,18 +653,21 @@ class Task(GraphNode):
         self.func = func
         if isinstance(func, Task):
             raise TypeError("Cannot nest tasks")
-        self.args = tuple(
-            Alias(obj.key) if isinstance(obj, TaskRef) else obj for obj in args
-        )
-        self.kwargs = {
-            k: Alias(v.key) if isinstance(v, TaskRef) else v for k, v in kwargs.items()
-        }
+
+        self.args = args
+        self.kwargs = kwargs
         _dependencies: set[KeyType] | None = None
-        for a in itertools.chain(self.args, self.kwargs.values()):
-            if isinstance(a, GraphNode):
+        for a in itertools.chain(args, kwargs.values()):
+            if isinstance(a, TaskRef):
                 if _dependencies is None:
-                    _dependencies = set()
-                _dependencies.update(a.dependencies)
+                    _dependencies = {a.key}
+                else:
+                    _dependencies.add(a.key)
+            elif isinstance(a, GraphNode) and a.dependencies:
+                if _dependencies is None:
+                    _dependencies = set(a.dependencies)
+                else:
+                    _dependencies.update(a.dependencies)
         if _dependencies:
             self._dependencies = frozenset(_dependencies)
         else:
@@ -739,19 +748,18 @@ class Task(GraphNode):
 
     def __call__(self, values=()):
         self._verify_values(values)
-        new_argspec = tuple(
-            a({k: values[k] for k in a.dependencies}) if isinstance(a, GraphNode) else a
-            for a in self.args
-        )
+
+        def _eval(a):
+            if isinstance(a, GraphNode):
+                return a({k: values[k] for k in a.dependencies})
+            elif isinstance(a, TaskRef):
+                return values[a.key]
+            else:
+                return a
+
+        new_argspec = tuple(map(_eval, self.args))
         if self.kwargs:
-            kwargs = {
-                k: (
-                    kw({k: values[k] for k in kw.dependencies})
-                    if isinstance(kw, GraphNode)
-                    else kw
-                )
-                for k, kw in self.kwargs.items()
-            }
+            kwargs = {k: _eval(kw) for k, kw in self.kwargs.items()}
             return self.func(*new_argspec, **kwargs)
         return self.func(*new_argspec)
 
@@ -788,11 +796,19 @@ class Task(GraphNode):
         }
         if subs_filtered:
             new_args = tuple(
-                a.substitute(subs_filtered) if isinstance(a, GraphNode) else a
+                (
+                    a.substitute(subs_filtered)
+                    if isinstance(a, (GraphNode, TaskRef))
+                    else a
+                )
                 for a in self.args
             )
             new_kwargs = {
-                k: v.substitute(subs_filtered) if isinstance(v, GraphNode) else v
+                k: (
+                    v.substitute(subs_filtered)
+                    if isinstance(v, (GraphNode, TaskRef))
+                    else v
+                )
                 for k, v in self.kwargs.items()
             }
             return type(self)(
@@ -866,7 +882,11 @@ class NestedContainer(Task, Iterable):
             return self
         return type(self)(
             *(
-                a.substitute(subs_filtered) if isinstance(a, GraphNode) else a
+                (
+                    a.substitute(subs_filtered)
+                    if isinstance(a, (GraphNode, TaskRef))
+                    else a
+                )
                 for a in self.args
             )
         )
@@ -946,7 +966,9 @@ class Dict(NestedContainer, Mapping):
         new_args = []
         for arg in self.args:
             new_arg = (
-                arg.substitute(subs_filtered) if isinstance(arg, GraphNode) else arg
+                arg.substitute(subs_filtered)
+                if isinstance(arg, (GraphNode, TaskRef))
+                else arg
             )
             new_args.append(new_arg)
         return type(self)(new_args)
diff --git a/dask/array/_shuffle.py b/dask/array/_shuffle.py
index 5851a4a99..5ece62851 100644
--- a/dask/array/_shuffle.py
+++ b/dask/array/_shuffle.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import copy
 import math
 from functools import reduce
 from itertools import count, product
@@ -192,8 +191,6 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
         else:
             return chunks, {}
 
-    indexer = copy.deepcopy(indexer)
-
     chunksize_tolerance = config.get("array.chunk-size-tolerance")
     chunk_size_limit = int(sum(chunks[axis]) / len(chunks[axis]) * chunksize_tolerance)
 
@@ -231,13 +228,10 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
         old_index: (in_name,) + old_index
         for old_index in np.ndindex(tuple([len(c) for c in chunks]))
     }
-
     for new_chunk_idx, new_chunk_taker in enumerate(new_chunks):
         new_chunk_taker = np.array(new_chunk_taker)
         sorter = np.argsort(new_chunk_taker).astype(dtype)
-        sorter_key = sorter_name + tokenize(sorter)
-        # low level fusion can't deal with arrays on first position
-        merges[sorter_key] = DataNode(sorter_key, (1, sorter))
+        sorter_key = None
 
         sorted_array = new_chunk_taker[sorter]
         source_chunk_nr, taker_boundary = np.unique(
@@ -286,6 +280,10 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
             merge_suffix = convert_key(chunk_tuple, new_chunk_idx, axis)
             out_name_merge = (out_name,) + merge_suffix
             if len(merge_keys) > 1:
+                if sorter_key is None:
+                    sorter_key = sorter_name + tokenize(sorter)
+                    # low level fusion can't deal with arrays on first position
+                    merges[sorter_key] = DataNode(sorter_key, (1, sorter))
                 merges[out_name_merge] = Task(
                     out_name_merge,
                     concatenate_arrays,
diff --git a/dask/array/core.py b/dask/array/core.py
index 5123c4ce3..3f0215ed0 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -45,7 +45,7 @@ from dask.array.dispatch import (  # noqa: F401
     tensordot_lookup,
 )
 from dask.array.numpy_compat import NUMPY_GE_200, _Recurser
-from dask.array.slicing import replace_ellipsis, setitem_array, slice_array
+from dask.array.slicing import SlicingNoop, replace_ellipsis, setitem_array, slice_array
 from dask.array.utils import (
     asanyarray_safe,
     asarray_safe,
@@ -1934,7 +1934,14 @@ class Array(DaskMethodsMixin):
         if isinstance(key, Array):
             from dask.array.routines import where
 
-            if key.shape != self.shape:
+            left_shape = np.array(key.shape)
+            right_shape = np.array(self.shape)
+
+            # We want to treat unknown shape on *either* sides as a match
+            match = left_shape == right_shape
+            match |= np.isnan(left_shape) | np.isnan(right_shape)
+
+            if not match.all():
                 raise IndexError(
                     f"boolean index shape {key.shape} must match indexed array's "
                     f"{self.shape}."
@@ -2028,7 +2035,10 @@ class Array(DaskMethodsMixin):
             return self
 
         out = "getitem-" + tokenize(self, index2)
-        dsk, chunks = slice_array(out, self.name, self.chunks, index2, self.itemsize)
+        try:
+            dsk, chunks = slice_array(out, self.name, self.chunks, index2)
+        except SlicingNoop:
+            return self
 
         graph = HighLevelGraph.from_collections(out, dsk, dependencies=[self])
 
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
index 46d8ca673..76419dc17 100644
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -16,7 +16,7 @@ from dask._task_spec import Alias, DataNode, Task, TaskRef
 from dask.array.chunk import getitem
 from dask.base import is_dask_collection, tokenize
 from dask.highlevelgraph import HighLevelGraph
-from dask.utils import _deprecated, cached_cumsum, is_arraylike
+from dask.utils import _deprecated, cached_cumsum, disable_gc, is_arraylike
 
 colon = slice(None, None, None)
 
@@ -97,7 +97,8 @@ def sanitize_index(ind):
         raise TypeError("Invalid index type", type(ind), ind)
 
 
-def slice_array(out_name, in_name, blockdims, index, itemsize):
+@disable_gc()
+def slice_array(out_name, in_name, blockdims, index):
     """
     Main function for array slicing
 
@@ -118,8 +119,6 @@ def slice_array(out_name, in_name, blockdims, index, itemsize):
       This is the dask variable output name
     blockshape - iterable of integers
     index - iterable of integers, slices, lists, or None
-    itemsize : int
-        The number of bytes required for each element of the array.
 
     Returns
     -------
@@ -140,10 +139,10 @@ def slice_array(out_name, in_name, blockdims, index, itemsize):
     --------
     >>> from pprint import pprint
     >>> dsk, blockdims = slice_array('y', 'x', [(20, 20, 20, 20, 20)],
-    ...                              (slice(10, 35),), 8)
+    ...                              (slice(10, 35),))
     >>> pprint(dsk)  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
-    {('y', 0): <Task ('y', 0) getitem(Alias(('x', 0)), (slice(10, 20, 1),))>,
-     ('y', 1): <Task ('y', 1) getitem(Alias(('x', 1)), (slice(0, 15, 1),))>}
+    {('y', 0): <Task ('y', 0) getitem(TaskRef(('x', 0)), (slice(10, 20, 1),))>,
+     ('y', 1): <Task ('y', 1) getitem(TaskRef(('x', 1)), (slice(0, 15, 1),))>}
     >>> blockdims
     ((10, 15),)
 
@@ -162,11 +161,7 @@ def slice_array(out_name, in_name, blockdims, index, itemsize):
     if all(
         isinstance(index, slice) and index == slice(None, None, None) for index in index
     ):
-        suffixes = product(*[range(len(bd)) for bd in blockdims])
-        dsk = {
-            (out_name,) + s: Alias((out_name,) + s, (in_name,) + s) for s in suffixes
-        }
-        return dsk, blockdims
+        raise SlicingNoop()
 
     # Add in missing colons at the end as needed.  x[5] -> x[5, :, :]
     not_none_count = sum(i is not None for i in index)
@@ -209,7 +204,7 @@ def slice_with_newaxes(out_name, in_name, blockdims, index):
         for k, v in dsk.items():
             if k[0] == out_name:
                 k2 = (out_name,) + expand(k[1:], 0)
-                if isinstance(v.args[1], Alias):
+                if isinstance(v.args[1], TaskRef):
                     # positional indexing with newaxis
                     indexer = expand_orig(dsk[v.args[1].key].value[1], None)
                     tok = "shuffle-taker-" + tokenize(indexer)
@@ -563,6 +558,12 @@ def issorted(seq):
     return np.all(seq[:-1] <= seq[1:])
 
 
+class SlicingNoop(Exception):
+    """This indicates that a slicing operation is a no-op. The caller has to handle this"""
+
+    pass
+
+
 def take(outname, inname, chunks, index, axis=0):
     """Index array with an iterable of index
 
@@ -581,9 +582,6 @@ def take(outname, inname, chunks, index, axis=0):
     >>> chunks
     ((4,),)
 
-    When any indexed blocks would otherwise grow larger than
-    dask.config.array.chunk-size, we will split them to avoid
-    growing chunksizes.
     """
 
     if not np.isnan(chunks[axis]).any():
@@ -592,14 +590,7 @@ def take(outname, inname, chunks, index, axis=0):
 
         arange = arange_safe(np.sum(chunks[axis]), like=index)
         if len(index) == len(arange) and np.abs(index - arange).sum() == 0:
-            # TODO: This should be a real no-op, but the call stack is
-            # too deep to do this efficiently for now
-            chunk_tuples = list(product(*(range(len(c)) for i, c in enumerate(chunks))))
-            graph = {
-                (outname,) + c: Alias((outname,) + c, (inname,) + c)
-                for c in chunk_tuples
-            }
-            return tuple(chunks), graph
+            raise SlicingNoop()
 
         average_chunk_size = int(sum(chunks[axis]) / len(chunks[axis]))
 
@@ -619,7 +610,7 @@ def take(outname, inname, chunks, index, axis=0):
         slices = [slice(None)] * len(chunks)
         slices[axis] = list(index)
         slices = tuple(slices)
-        chunk_tuples = list(product(*(range(len(c)) for i, c in enumerate(chunks))))
+        chunk_tuples = product(*(range(len(c)) for i, c in enumerate(chunks)))
         dsk = {
             (outname,)
             + ct: Task((outname,) + ct, getitem, TaskRef((inname,) + ct), slices)
diff --git a/dask/blockwise.py b/dask/blockwise.py
index eea650a9d..4b0d11f85 100644
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -271,10 +271,10 @@ def blockwise(
     ...         numblocks={'x': (2, 2)}
     ...     )
     ... )  # doctest: +NORMALIZE_WHITESPACE
-    {('z', 0, 0): <Task ('z', 0, 0) inc(Alias(('x', 0, 0)))>,
-    ('z', 0, 1): <Task ('z', 0, 1) inc(Alias(('x', 0, 1)))>,
-    ('z', 1, 0): <Task ('z', 1, 0) inc(Alias(('x', 1, 0)))>,
-    ('z', 1, 1): <Task ('z', 1, 1) inc(Alias(('x', 1, 1)))>}
+    {('z', 0, 0): <Task ('z', 0, 0) inc(TaskRef(('x', 0, 0)))>,
+    ('z', 0, 1): <Task ('z', 0, 1) inc(TaskRef(('x', 0, 1)))>,
+    ('z', 1, 0): <Task ('z', 1, 0) inc(TaskRef(('x', 1, 0)))>,
+    ('z', 1, 1): <Task ('z', 1, 1) inc(TaskRef(('x', 1, 1)))>}
 
     Simple operation on two datasets x and y
 
@@ -289,10 +289,10 @@ def blockwise(
     ...         numblocks={'x': (2, 2), 'y': (2, 2)}
     ...     )
     ... )  # doctest: +NORMALIZE_WHITESPACE
-    {('z', 0, 0): <Task ('z', 0, 0) add(Alias(('x', 0, 0)), Alias(('y', 0, 0)))>,
-    ('z', 0, 1): <Task ('z', 0, 1) add(Alias(('x', 0, 1)), Alias(('y', 0, 1)))>,
-    ('z', 1, 0): <Task ('z', 1, 0) add(Alias(('x', 1, 0)), Alias(('y', 1, 0)))>,
-    ('z', 1, 1): <Task ('z', 1, 1) add(Alias(('x', 1, 1)), Alias(('y', 1, 1)))>}
+    {('z', 0, 0): <Task ('z', 0, 0) add(TaskRef(('x', 0, 0)), TaskRef(('y', 0, 0)))>,
+    ('z', 0, 1): <Task ('z', 0, 1) add(TaskRef(('x', 0, 1)), TaskRef(('y', 0, 1)))>,
+    ('z', 1, 0): <Task ('z', 1, 0) add(TaskRef(('x', 1, 0)), TaskRef(('y', 1, 0)))>,
+    ('z', 1, 1): <Task ('z', 1, 1) add(TaskRef(('x', 1, 1)), TaskRef(('y', 1, 1)))>}
 
     Operation that flips one of the datasets
 
@@ -307,10 +307,10 @@ def blockwise(
     ...         numblocks={'x': (2, 2), 'y': (2, 2)}
     ...     )
     ... )  # doctest: +NORMALIZE_WHITESPACE
-    {('z', 0, 0): <Task ('z', 0, 0) addT(Alias(('x', 0, 0)), Alias(('y', 0, 0)))>,
-    ('z', 0, 1): <Task ('z', 0, 1) addT(Alias(('x', 0, 1)), Alias(('y', 1, 0)))>,
-    ('z', 1, 0): <Task ('z', 1, 0) addT(Alias(('x', 1, 0)), Alias(('y', 0, 1)))>,
-    ('z', 1, 1): <Task ('z', 1, 1) addT(Alias(('x', 1, 1)), Alias(('y', 1, 1)))>}
+    {('z', 0, 0): <Task ('z', 0, 0) addT(TaskRef(('x', 0, 0)), TaskRef(('y', 0, 0)))>,
+    ('z', 0, 1): <Task ('z', 0, 1) addT(TaskRef(('x', 0, 1)), TaskRef(('y', 1, 0)))>,
+    ('z', 1, 0): <Task ('z', 1, 0) addT(TaskRef(('x', 1, 0)), TaskRef(('y', 0, 1)))>,
+    ('z', 1, 1): <Task ('z', 1, 1) addT(TaskRef(('x', 1, 1)), TaskRef(('y', 1, 1)))>}
 
     Dot product with contraction over ``j`` index.  Yields list arguments
 
@@ -326,24 +326,25 @@ def blockwise(
     ... )  # doctest: +SKIP
     {('z', 0,  0):
         <Task ('z', 0, 0) dotmany(
-        List((Alias(('x', 0, 0)), Alias(('x', 0, 1)))),
-        List((Alias(('y', 0, 0)), Alias(('y', 1, 0)))))
+            List((TaskRef(('x', 0, 0)), TaskRef(('x', 0, 1)))),
+            List((TaskRef(('y', 0, 0)), TaskRef(('y', 1, 0)))))
         >,
     ('z', 0,  1):
         <Task ('z', 0, 1) dotmany(
-        List((Alias(('x', 0, 0)), Alias(('x', 0, 1)))),
-        List((Alias(('y', 0, 1)), Alias(('y', 1, 1)))))
+            List((TaskRef(('x', 0, 0)), TaskRef(('x', 0, 1)))),
+            List((TaskRef(('y', 0, 1)), TaskRef(('y', 1, 1)))))
         >,
     ('z', 1,  0):
         <Task ('z', 1, 0) dotmany(
-        List((Alias(('x', 1, 0)), Alias(('x', 1, 1)))),
-        List((Alias(('y', 0, 0)), Alias(('y', 1, 0)))))
+            List((TaskRef(('x', 1, 0)), TaskRef(('x', 1, 1)))),
+            List((TaskRef(('y', 0, 0)), TaskRef(('y', 1, 0)))
+        )
         >,
-    ('z', 1,  1):
+    ('z', 1, 1):
         <Task ('z', 1, 1) dotmany(
-        List((Alias(('x', 1, 0)), Alias(('x', 1, 1)))),
-        List((Alias(('y', 0, 1)), Alias(('y', 1, 1)))))
-        >
+            List((TaskRef(('x', 1, 0)), TaskRef(('x', 1, 1)))),
+            List((TaskRef(('y', 0, 1)), TaskRef(('y', 1, 1)))
+        )
     }
 
     Pass ``concatenate=True`` to concatenate arrays ahead of time.
@@ -363,10 +364,10 @@ def blockwise(
     ...     )
     ... )  # doctest: +SKIP
     {
-    ('z', 0, 0): <Task ('z', 0, 0) add(Alias(('x', 0, 0)), Alias(('y', 0, 0)))>,
-    ('z', 0, 1): <Task ('z', 0, 1) add(Alias(('x', 0, 1)), Alias(('y', 0, 1)))>,
-    ('z', 1, 0): <Task ('z', 1, 0) add(Alias(('x', 0, 0)), Alias(('y', 1, 0)))>,
-    ('z', 1, 1): <Task ('z', 1, 1) add(Alias(('x', 0, 1)), Alias(('y', 1, 1)))>
+    ('z', 0, 0): <Task ('z', 0, 0) add(TaskRef(('x', 0, 0)), TaskRef(('y', 0, 0)))>,
+    ('z', 0, 1): <Task ('z', 0, 1) add(TaskRef(('x', 0, 1)), TaskRef(('y', 0, 1)))>,
+    ('z', 1, 0): <Task ('z', 1, 0) add(TaskRef(('x', 0, 0)), TaskRef(('y', 1, 0)))>,
+    ('z', 1, 1): <Task ('z', 1, 1) add(TaskRef(('x', 0, 1)), TaskRef(('y', 1, 1)))>
     }
 
     Include literals by indexing with ``None``
@@ -380,8 +381,8 @@ def blockwise(
     ...         numblocks={'x': (2, )}
     ...     )
     ... )  # doctest: +NORMALIZE_WHITESPACE
-    {('z', 0): <Task ('z', 0) add(Alias(('x', 0)), DataNode(100))>,
-    ('z', 1): <Task ('z', 1) add(Alias(('x', 1)), DataNode(100))>}
+    {('z', 0): <Task ('z', 0) add(TaskRef(('x', 0)), DataNode(100))>,
+    ('z', 1): <Task ('z', 1) add(TaskRef(('x', 1)), DataNode(100))>}
 
     If the broadcasted value is a delayed object or other dask collection, the
     key has to be wrapped appropriately as an Alias object.
@@ -398,8 +399,8 @@ def blockwise(
     ...         numblocks={'x': (2, )}
     ...     )
     ... )  # doctest: +NORMALIZE_WHITESPACE
-    {('z', 0): <Task ('z', 0) add(Alias(('x', 0)), Alias('dinc'))>,
-    ('z', 1): <Task ('z', 1) add(Alias(('x', 1)), Alias('dinc'))>}
+    {('z', 0): <Task ('z', 0) add(TaskRef(('x', 0)), Alias('dinc'))>,
+    ('z', 1): <Task ('z', 1) add(TaskRef(('x', 1)), Alias('dinc'))>}
 
     See Also
     --------
diff --git a/dask/delayed.py b/dask/delayed.py
index 6a5482053..114da263f 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -153,13 +153,13 @@ def unpack_collections(expr, _return_collections=True):
     >>> b = delayed(2, 'b')
     >>> task, collections = unpack_collections([a, b, 3])
     >>> task
-    List((Alias('a'), Alias('b'), 3))
+    List((TaskRef('a'), TaskRef('b'), 3))
     >>> collections
     (Delayed('a'), Delayed('b'))
 
     >>> task, collections = unpack_collections({a: 1, b: 2})
     >>> task
-    Dict(Alias('a'): 1, Alias('b'): 2)
+    Dict(a: 1, b: 2)
     >>> collections
     (Delayed('a'), Delayed('b'))
     """
diff --git a/dask/order.py b/dask/order.py
index b6db4578b..e33489029 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -195,7 +195,7 @@ def order(
     runnable_hull = set()
     reachable_hull = set()
 
-    runnable = []
+    runnable: list[Key] = []
 
     known_runnable_paths: dict[Key, list[list[Key]]] = {}
     known_runnable_paths_pop = known_runnable_paths.pop
@@ -203,12 +203,8 @@ def order(
     crit_path_counter = 0
     scrit_path: set[Key] = set()
     _crit_path_counter_offset: int | float = 0
-
     _sort_keys_cache: dict[Key, tuple[int, int, int, int, str]] = {}
 
-    leafs_connected_to_loaded_roots: set[Key] = set()
-    processed_roots = set()
-
     def sort_key(x: Key) -> tuple[int, int, int, int, str]:
         try:
             return _sort_keys_cache[x]
@@ -231,6 +227,7 @@ def order(
         # recursion depth errors. This is the only reason for the while loop
         next_items = [item]
         nonlocal i
+        nonlocal min_leaf_degree  # type: ignore[misc]
         while next_items:
             item = next_items.pop()
             runnable_hull.discard(item)
@@ -248,7 +245,25 @@ def order(
             if item not in external_keys:
                 i += 1
             if item in root_nodes:
-                processed_roots.add(item)
+                for leaf in leafs_connected[item]:
+                    if leaf in leaf_nodes:
+                        degree = leafs_degree.pop(leaf)
+                        leafs_per_degree[degree].remove(leaf)
+
+                        new_degree = degree - 1
+                        if new_degree > 0:
+                            if new_degree < min_leaf_degree:
+                                min_leaf_degree = new_degree
+                            leafs_per_degree[new_degree].add(leaf)
+                            leafs_degree[leaf] = new_degree
+                        elif not leafs_per_degree[degree]:
+                            assert degree == min_leaf_degree
+                            while min_leaf_degree != max_leaf_degree and (
+                                min_leaf_degree not in leafs_per_degree
+                                or not leafs_per_degree[min_leaf_degree]
+                            ):
+                                min_leaf_degree += 1
+
             # Note: This is a `set` and therefore this introduces a certain
             # randomness. However, this randomness should not have any impact on
             # the final result since the `process_runnable` should produce
@@ -396,42 +411,38 @@ def order(
     # writing, the most expensive part of ordering is the prep work (mostly
     # connected roots + sort_key) which can be reused for multiple orderings.
 
-    def get_target(longest_path: bool = False) -> Key:
-        # Some topologies benefit if the node with the most dependencies
-        # is used as first choice, others benefit from the opposite.
-        candidates = leaf_nodes
-        skey: Callable = sort_key
-        if reachable_hull:
-            skey = lambda k: (num_needed[k], sort_key(k))
-
-        all_leafs_accessible = len(leafs_connected_to_loaded_roots) >= len(candidates)
-        if reachable_hull and not all_leafs_accessible:
-            # Avoid this branch if we can already access all leafs. Just pick
-            # one of them without the expensive selection process in this case.
-
-            candidates = reachable_hull & candidates
+    # Degree in this context is the number root nodes that have to be loaded for
+    # this leaf to become accessible. Zero means the leaf is already accessible
+    # in which case it _should_ either already be in result or be accessible vie
+    # process_runnables
+    # When picking a new target, we prefer the leafs with the least number of
+    # roots that need loading.
+    leafs_degree = {}
+    leafs_per_degree = defaultdict(set)
+    min_leaf_degree = len(dsk)
+    max_leaf_degree = len(dsk)
+    for leaf in leaf_nodes - root_nodes:
+        degree = len(roots_connected[leaf])
+        min_leaf_degree = min(min_leaf_degree, degree)
+        max_leaf_degree = max(max_leaf_degree, degree)
+        leafs_degree[leaf] = degree
+        leafs_per_degree[degree].add(leaf)
+
+    def get_target() -> Key:
+        # If we're already mid run and there is a runnable_hull we'll attempt to
+        # pick the next target in a way that minimizes the number of additional
+        # root nodes that are needed
+        all_leafs_accessible = min_leaf_degree == max_leaf_degree
+        is_trivial_lookup = not reachable_hull or all_leafs_accessible
+        if not is_trivial_lookup:
+            candidates = reachable_hull & leafs_per_degree[min_leaf_degree]
             if not candidates:
-                # We can't reach a leaf node directly, but we still have nodes
-                # with results in memory, these notes can inform our path towards
-                # a new preferred leaf node.
-                for r in processed_roots:
-                    leafs_connected_to_loaded_roots.update(leafs_connected[r])
-                processed_roots.clear()
-
-                leafs_connected_to_loaded_roots.intersection_update(leaf_nodes)
-                candidates = leafs_connected_to_loaded_roots
-            else:
-                leafs_connected_to_loaded_roots.update(candidates)
-
-        elif not reachable_hull:
-            # We reach a new and independent branch, so throw away previous branch
-            leafs_connected_to_loaded_roots.clear()
-
-        if longest_path and (not reachable_hull or all_leafs_accessible):
-            return leaf_nodes_sorted.pop()
+                candidates = leafs_per_degree[min_leaf_degree]
+            # Even without reachable hull overlap this should be relatively
+            # small so one full pass should be fine
+            return min(candidates, key=sort_key)
         else:
-            # FIXME: This can be very expensive
-            return min(candidates, key=skey)
+            return leaf_nodes_sorted.pop()
 
     def use_longest_path() -> bool:
         size = 0
@@ -450,10 +461,10 @@ def order(
 
         return True
 
+    # Some topologies benefit if the node with the most dependencies
+    # is used as first choice, others benefit from the opposite.
     longest_path = use_longest_path()
-    leaf_nodes_sorted = []
-    if longest_path:
-        leaf_nodes_sorted = sorted(leaf_nodes, key=sort_key, reverse=False)
+    leaf_nodes_sorted = sorted(leaf_nodes, key=sort_key, reverse=not longest_path)
 
     # *************************************************************************
     # CORE ALGORITHM STARTS HERE
@@ -553,7 +564,7 @@ def order(
         assert not scrit_path
 
         # A. Build the critical path
-        target = get_target(longest_path=longest_path)
+        target = get_target()
         next_deps = dependencies[target]
         path_append(target)
 
diff --git a/dask/utils.py b/dask/utils.py
index af99912d3..ee9226ecf 100644
--- a/dask/utils.py
+++ b/dask/utils.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import codecs
 import functools
+import gc
 import inspect
 import os
 import re
@@ -12,7 +13,7 @@ import types
 import uuid
 import warnings
 from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Set
-from contextlib import contextmanager, nullcontext, suppress
+from contextlib import ContextDecorator, contextmanager, nullcontext, suppress
 from datetime import datetime, timedelta
 from errno import ENOENT
 from functools import wraps
@@ -742,6 +743,8 @@ class Dispatch:
     def dispatch(self, cls):
         """Return the function implementation for the given ``cls``"""
         lk = self._lookup
+        if cls in lk:
+            return lk[cls]
         for cls2 in cls.__mro__:
             # Is a lazy registration function present?
             toplevel, _, _ = cls2.__module__.partition(".")
@@ -752,7 +755,10 @@ class Dispatch:
             else:
                 register()
                 self._lazy.pop(toplevel, None)
-                return self.dispatch(cls)  # recurse
+                meth = self.dispatch(cls)  # recurse
+                lk[cls] = meth
+                lk[cls2] = meth
+                return meth
             try:
                 impl = lk[cls2]
             except KeyError:
@@ -2299,3 +2305,20 @@ def unzip(ls, nout):
     if not out:
         out = [()] * nout
     return out
+
+
+class disable_gc(ContextDecorator):
+    """Context manager to disable garbage collection."""
+
+    def __init__(self, collect=False):
+        self.collect = collect
+        self._gc_enabled = gc.isenabled()
+
+    def __enter__(self):
+        gc.disable()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self._gc_enabled:
+            gc.enable()
+        return False
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index 849301c25..5380ef914 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -5,6 +5,16 @@ Changelog
 
     This is not exhaustive. For an exhaustive list of changes, see the git log.
 
+
+.. _v2025.5.0:
+
+2025.5.0
+--------
+
+- Fixed Array setitem when both the array and the indexer have unknown shape (:issue:`11753`) `Tom Augspurger`_
+
+
+
 .. _v2025.4.1:
 
 2025.4.1


Test patch:
diff --git a/dask/array/tests/test_slicing.py b/dask/array/tests/test_slicing.py
index 28eef7841..daec4c1fd 100644
--- a/dask/array/tests/test_slicing.py
+++ b/dask/array/tests/test_slicing.py
@@ -6,14 +6,15 @@ import warnings
 import pytest
 
 from dask._task_spec import Alias, Task, TaskRef
+from dask.delayed import delayed
 
 np = pytest.importorskip("numpy")
 
 import dask
 import dask.array as da
-from dask import config
 from dask.array.chunk import getitem
 from dask.array.slicing import (
+    SlicingNoop,
     _sanitize_index_element,
     _slice_1d,
     make_block_sorted_slices,
@@ -184,7 +185,7 @@ def test_slice_array_1d():
         ("y", 2): Task(("y", 2), getitem, TaskRef(("x", 2)), (slice(0, 25, 2),)),
         ("y", 3): Task(("y", 3), getitem, TaskRef(("x", 3)), (slice(1, 25, 2),)),
     }
-    result, chunks = slice_array("y", "x", [[25] * 4], [slice(24, None, 2)], 8)
+    result, chunks = slice_array("y", "x", [[25] * 4], [slice(24, None, 2)])
 
     assert expected == result
 
@@ -195,7 +196,7 @@ def test_slice_array_1d():
         ("y", 2): Task(("y", 2), getitem, TaskRef(("x", 3)), (slice(1, 25, 2),)),
     }
 
-    result, chunks = slice_array("y", "x", [[25] * 4], [slice(26, None, 2)], 8)
+    result, chunks = slice_array("y", "x", [[25] * 4], [slice(26, None, 2)])
     assert expected == result
 
     # x[24::2]
@@ -205,7 +206,7 @@ def test_slice_array_1d():
         ("y", 2): Task(("y", 2), getitem, TaskRef(("x", 2)), (slice(0, 25, 2),)),
         ("y", 3): Task(("y", 3), getitem, TaskRef(("x", 3)), (slice(1, 25, 2),)),
     }
-    result, chunks = slice_array("y", "x", [(25,) * 4], (slice(24, None, 2),), 8)
+    result, chunks = slice_array("y", "x", [(25,) * 4], (slice(24, None, 2),))
 
     assert expected == result
 
@@ -216,7 +217,7 @@ def test_slice_array_1d():
         ("y", 2): Task(("y", 2), getitem, TaskRef(("x", 3)), (slice(1, 25, 2),)),
     }
 
-    result, chunks = slice_array("y", "x", [(25,) * 4], (slice(26, None, 2),), 8)
+    result, chunks = slice_array("y", "x", [(25,) * 4], (slice(26, None, 2),))
     assert expected == result
 
 
@@ -242,7 +243,6 @@ def test_slice_array_2d():
         "x",
         [[20], [20, 20, 5]],
         [slice(13, None, 2), slice(10, None, 1)],
-        itemsize=8,
     )
 
     assert expected == result
@@ -254,39 +254,28 @@ def test_slice_array_2d():
         ("y", 2): Task(("y", 2), getitem, TaskRef(("x", 0, 2)), (5, slice(None))),
     }
 
-    result, chunks = slice_array(
-        "y", "x", ([20], [20, 20, 5]), [5, slice(10, None, 1)], 8
-    )
+    result, chunks = slice_array("y", "x", ([20], [20, 20, 5]), [5, slice(10, None, 1)])
 
     assert expected == result
 
 
 def test_slice_optimizations():
     # bar[:]
-    expected = {("foo", 0): Alias(("foo", 0), ("bar", 0))}
-    result, chunks = slice_array("foo", "bar", [[100]], (slice(None, None, None),), 8)
-    assert expected == result
+    with pytest.raises(SlicingNoop):
+        slice_array("foo", "bar", [[100]], (slice(None, None, None),))
 
     # bar[:,:,:]
-    expected = {
-        ("foo", 0): Alias(("foo", 0), ("bar", 0)),
-        ("foo", 1): Alias(("foo", 1), ("bar", 1)),
-        ("foo", 2): Alias(("foo", 2), ("bar", 2)),
-    }
-    result, chunks = slice_array(
-        "foo",
-        "bar",
-        [(100, 1000, 10000)],
-        (slice(None, None, None), slice(None, None, None), slice(None, None, None)),
-        itemsize=8,
-    )
-    assert expected == result
+    with pytest.raises(SlicingNoop):
+        slice_array(
+            "foo",
+            "bar",
+            [(100, 1000, 10000)],
+            (slice(None, None, None), slice(None, None, None), slice(None, None, None)),
+        )
 
 
 def test_slicing_with_singleton_indices():
-    result, chunks = slice_array(
-        "y", "x", ([5, 5], [5, 5]), (slice(0, 5), 8), itemsize=8
-    )
+    result, chunks = slice_array("y", "x", ([5, 5], [5, 5]), (slice(0, 5), 8))
 
     expected = {
         ("y", 0): Task(
@@ -303,7 +292,6 @@ def test_slicing_with_newaxis():
         "x",
         ([5, 5], [5, 5]),
         (slice(0, 3), None, slice(None, None, None)),
-        itemsize=8,
     )
 
     expected = {
@@ -339,26 +327,20 @@ def test_take_sorted():
     chunks, dsk = take("y-y", "x", [(20, 20, 20, 20)], [1, 3, 5, 47], axis=0)
     assert len(dsk) == 6
     assert chunks == ((4,),)
-
-    chunks, dsk = take("y", "x", [(20, 20, 20, 20)], np.arange(0, 80), axis=0)
-    assert len(dsk) == 4
-    assert chunks == ((20, 20, 20, 20),)
+    with pytest.raises(SlicingNoop):
+        take("y", "x", [(20, 20, 20, 20)], np.arange(0, 80), axis=0)
 
 
 def test_slicing_chunks():
-    result, chunks = slice_array(
-        "y", "x", ([5, 5], [5, 5]), (1, np.array([2, 0, 3])), itemsize=8
-    )
+    result, chunks = slice_array("y", "x", ([5, 5], [5, 5]), (1, np.array([2, 0, 3])))
     assert chunks == ((3,),)
 
     result, chunks = slice_array(
-        "y", "x", ([5, 5], [5, 5]), (slice(0, 7), np.array([2, 0, 3])), itemsize=8
+        "y", "x", ([5, 5], [5, 5]), (slice(0, 7), np.array([2, 0, 3]))
     )
     assert chunks == ((5, 2), (3,))
 
-    result, chunks = slice_array(
-        "y", "x", ([5, 5], [5, 5]), (slice(0, 7), 1), itemsize=8
-    )
+    result, chunks = slice_array("y", "x", ([5, 5], [5, 5]), (slice(0, 7), 1))
     assert chunks == ((5, 2),)
 
 
@@ -368,14 +350,12 @@ def test_slicing_with_numpy_arrays():
         "x",
         ((3, 3, 3, 1), (3, 3, 3, 1)),
         (np.array([1, 2, 9]), slice(None, None, None)),
-        itemsize=8,
     )
     b, bd2 = slice_array(
         "y-y",
         "x",
         ((3, 3, 3, 1), (3, 3, 3, 1)),
         (np.array([1, 2, 9]), slice(None, None, None)),
-        itemsize=8,
     )
 
     assert bd1 == bd2
@@ -384,7 +364,7 @@ def test_slicing_with_numpy_arrays():
     i = [False, True, True, False, False, False, False, False, False, True]
     index = (i, slice(None, None, None))
     index = normalize_index(index, (10, 10))
-    c, bd3 = slice_array("y-y", "x", ((3, 3, 3, 1), (3, 3, 3, 1)), index, itemsize=8)
+    c, bd3 = slice_array("y-y", "x", ((3, 3, 3, 1), (3, 3, 3, 1)), index)
     assert bd1 == bd3
     np.testing.assert_equal(a, c)
 
@@ -574,7 +554,7 @@ def test_sanitize_index():
 def test_uneven_blockdims():
     blockdims = ((31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30), (100,))
     index = (slice(240, 270), slice(None))
-    dsk_out, bd_out = slice_array("in", "out", blockdims, index, itemsize=8)
+    dsk_out, bd_out = slice_array("in", "out", blockdims, index)
     sol = {
         ("in", 0, 0): Task(
             ("in", 0, 0),
@@ -594,7 +574,7 @@ def test_uneven_blockdims():
 
     blockdims = ((31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30),) * 2
     index = (slice(240, 270), slice(180, 230))
-    dsk_out, bd_out = slice_array("in", "out", blockdims, index, itemsize=8)
+    dsk_out, bd_out = slice_array("in", "out", blockdims, index)
     sol = {
         ("in", 0, 0): Task(
             ("in", 0, 0),
@@ -931,33 +911,23 @@ def test_take_avoids_large_chunks():
         index = np.array([0, 1] + [2] * 101 + [3])
         chunks2, dsk = take("a", "b", chunks, index)
         assert chunks2 == ((1,) * 104, (500,), (500,))
-        assert len(dsk) == 106
+        assert len(dsk) == 105
 
         index = np.array([0] * 101 + [1, 2, 3])
         chunks2, dsk = take("a", "b", chunks, index)
         assert chunks2 == ((1,) * 104, (500,), (500,))
-        assert len(dsk) == 106
+        assert len(dsk) == 105
 
         index = np.array([0, 1, 2] + [3] * 101)
         chunks2, dsk = take("a", "b", chunks, index)
         assert chunks2 == ((1,) * 104, (500,), (500,))
-        assert len(dsk) == 106
+        assert len(dsk) == 105
 
         chunks = ((500,), (1, 1, 1, 1), (500,))
         index = np.array([0, 1, 2] + [3] * 101)
         chunks2, dsk = take("a", "b", chunks, index, axis=1)
         assert chunks2 == ((500,), (1,) * 104, (500,))
-        assert len(dsk) == 106
-
-
-def test_take_uses_config():
-    with dask.config.set({"array.slicing.split-large-chunks": True}):
-        chunks = ((1, 1, 1, 1), (500,), (500,))
-        index = np.array([0, 1] + [2] * 101 + [3])
-        with config.set({"array.chunk-size": "10GB"}):
-            chunks2, dsk = take("a", "b", chunks, index)
-        assert chunks2 == ((1,) * 104, (500,), (500,))
-        assert len(dsk) == 106
+        assert len(dsk) == 105
 
 
 def test_pathological_unsorted_slicing():
@@ -1137,3 +1107,39 @@ def test_positional_indexer_newaxis():
     arr = da.array([0, 1, 2])
     new = arr[[True, True, False], np.newaxis]
     assert_eq(new, arr.compute()[[True, True, False], np.newaxis])
+
+
+@pytest.mark.parametrize(
+    "shapes",
+    [
+        (10, 10),
+        (np.nan, np.nan),
+        pytest.param(
+            (10, np.nan), marks=pytest.mark.xfail(reason="Not implemented", strict=True)
+        ),
+        pytest.param(
+            (np.nan, 10), marks=pytest.mark.xfail(reason="Not implemented", strict=True)
+        ),
+    ],
+)
+def test_boolean_mask_with_unknown_shape(
+    shapes: tuple[float | int, float | int],
+) -> None:
+    x_shape, mask_shape = shapes
+    arr = delayed(np.ones(10))
+    x = da.concatenate(
+        [
+            da.from_delayed(arr, shape=(x_shape,), dtype=float),
+            da.from_delayed(arr, shape=(x_shape,), dtype=float),
+        ]
+    )
+    mask = da.concatenate(
+        [
+            da.from_delayed(arr, shape=(mask_shape,), dtype=bool),
+            da.from_delayed(arr, shape=(mask_shape,), dtype=bool),
+        ]
+    )
+    x[mask] = 2
+
+    expected = np.full(20, 2.0)
+    assert_eq(x, expected)
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index 4074c5b78..b93c95eb0 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -560,10 +560,10 @@ def test_run_smaller_sections(abcde):
     }
     o = order(dsk)
     assert_topological_sort(dsk, o)
-    assert max(diagnostics(dsk)[1]) <= 4  # optimum is 3
-    # This is a mildly ambiguous example
+    # The actual order is a bit ambiguous but the optimal pressure is well
+    # defined, see also
     # https://github.com/dask/dask/pull/10535/files#r1337528255
-    assert (o[aa] < o[a] and o[dd] < o[a]) or (o[b] < o[e] and o[b] < o[cc])
+    assert max(diagnostics(dsk)[1]) == 3
 
 
 def test_local_parents_of_reduction(abcde):