Repository: dask/dask

Issue:
einsum needs more memory when chunking the tensor
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**Describe the issue**:
See the following code. First we make two random tensors ```A``` and ```B``` which's size are both ```(N,D,D,D,D)```, with chunksize```(N//nN,D//nD,D//nD,D//nD,D//nD)```. Then perform the ```einsum``` and we get the tensor ```C```. The shape of tensor ```C``` becomes ```(D,D,D,D,D,D,D,D,nN)``` at Layer3. Especially in the case of $\text{nN} \neq 1$, the memory usage increases $\text{nN}$ times
![image](https://user-images.githubusercontent.com/89069722/224502561-40fff3ef-9667-4642-aef7-518b751197bb.png)


**Minimal Complete Verifiable Example**:
```python
import dask.array as da
import opt_einsum as oe

N = 1024
D = 16

nN = 4
nD = 2

chunks=(N//nN,D//nD,D//nD,D//nD,D//nD)

A = da.random.random((N,D,D,D,D), chunks=chunks)
B = da.random.random((N,D,D,D,D), chunks=chunks)
C = da.einsum("aijkl,amnop->ijklmnop", A, B)
#C = oe.contract("aijkl,amnop->ijklmnop", A, B)
C.dask
```

**Anything else we need to know?**:

**Environment**:
- Dask version: 2023.1.1
- Python version: 3.10.9
- Operating System: Red Hat
- Install method (conda, pip, source): conda


Discussion:
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 48m 50s ⏱️ + 7m 18s
 13 241 tests + 1   12 179 ✅ + 1   1 062 💤 ±0  0 ❌ ±0 
138 138 runs  +11  119 049 ✅ +12  19 089 💤  - 1  0 ❌ ±0 

Results for commit 0f297bd9. ± Comparison against base commit a0783a8d.

[test-results]:data:application/gzip;base64,H4sIAA3eIGcC/02Myw6DIBQFf8Ww7oILVLn9GaM8ElKVBmFl+u9FpWhyNjMnmY1YN5mVvBrgj4asycWLdApDdH45uOM0q/zG82cC/tyvSaldMujwkm/32SVtWXV2cFN2tdSbEHwoJqTlTMu8wjUNSAVWWdJIZXW39MH3svLz7GIGQi3DbtQokGljLHLzFCjBWCnVMFqhBHCtWyDfH7sWXWYaAQAA


Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index e02e6cbb1..a79ceb35c 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -272,7 +272,10 @@ def convert_legacy_task(
     try:
         if isinstance(task, (bytes, int, float, str, tuple)):
             if task in all_keys:
-                return Alias(task)
+                if k is None:
+                    return Alias(task)
+                else:
+                    return Alias(k, target=task)
     except TypeError:
         # Unhashable
         pass
@@ -305,7 +308,10 @@ def convert_legacy_task(
         )
 
     elif isinstance(task, TaskRef):
-        return Alias(task.key)
+        if k is None:
+            return Alias(task.key)
+        else:
+            return Alias(k, target=task.key)
     else:
         return task
 
@@ -320,7 +326,13 @@ def convert_legacy_graph(
     new_dsk = {}
     for k, arg in dsk.items():
         t = convert_legacy_task(k, arg, all_keys, only_refs)
-        if not only_refs and isinstance(t, Alias) and t.key == k:
+        if (
+            not only_refs
+            and isinstance(t, Alias)
+            and isinstance(arg, TaskRef)
+            and t.key == arg.key
+        ):
+            # This detects cycles?
             continue
         new_dsk[k] = t
     new_dsk2 = {
@@ -379,6 +391,8 @@ def resolve_aliases(dsk: dict, keys: set, dependents: dict) -> dict:
                 if isinstance(tnew, Alias):
                     work.append(k)
                     seen.discard(k)
+                else:
+                    work.extend(tnew.dependencies)
 
         work.extend(t.dependencies)
     return dsk
@@ -653,7 +667,9 @@ class Task(GraphNode):
         )
 
     def __repr__(self) -> str:
-        if not self._repr:
+        # When `Task` is deserialized the constructor will not run and
+        # `self._repr` is thus undefined.
+        if not hasattr(self, "_repr") or not self._repr:
             self.unpack()
             head = funcname(self.func)
             tail = ")"
@@ -841,25 +857,34 @@ class DependenciesMapping(MutableMapping):
     def __init__(self, dsk):
         self.dsk = dsk
         self._removed = set()
+        # Set a copy of dsk to avoid dct resizing
+        self._cache = dsk.copy()
+        self._cache.clear()
 
     def __getitem__(self, key):
-        v = self.dsk[key]
-        if not isinstance(v, GraphNode):
-            from dask.core import get_dependencies
-
-            deps = get_dependencies(self.dsk, task=self.dsk[key])
+        if (val := self._cache.get(key)) is not None:
+            return val
         else:
-            deps = self.dsk[key].dependencies
-        if self._removed:
-            # deps is a frozenset but for good measure, let's not use -= since
-            # that _may_ perform an inplace mutation
-            deps = deps - self._removed
-        return deps
+            v = self.dsk[key]
+            try:
+                deps = v.dependencies
+            except AttributeError:
+                from dask.core import get_dependencies
+
+                deps = get_dependencies(self.dsk, task=v)
+
+            if self._removed:
+                # deps is a frozenset but for good measure, let's not use -= since
+                # that _may_ perform an inplace mutation
+                deps = deps - self._removed
+            self._cache[key] = deps
+            return deps
 
     def __iter__(self):
         return iter(self.dsk)
 
     def __delitem__(self, key: Any) -> None:
+        self._cache.clear()
         self._removed.add(key)
 
     def __setitem__(self, key: Any, value: Any) -> None:
diff --git a/dask/array/_shuffle.py b/dask/array/_shuffle.py
index 8299f4224..eedbb6aa0 100644
--- a/dask/array/_shuffle.py
+++ b/dask/array/_shuffle.py
@@ -92,22 +92,11 @@ def shuffle(x, indexer: list[list[int]], axis: int, chunks: Literal["auto"] = "a
     return Array(graph, out_name, chunks, meta=x)
 
 
-def _rechunk_other_dimensions(
-    x: Array, longest_group: int, axis: int, chunks: Literal["auto"]
-) -> Array:
-    assert chunks == "auto", "Only auto is supported for now"
-    chunksize_tolerance = config.get("array.chunk-size-tolerance")
-
-    if longest_group <= max(x.chunks[axis]) * chunksize_tolerance:
-        # We are staying below our threshold, so don't rechunk
-        return x
+def _calculate_new_chunksizes(
+    input_chunks, new_chunks, changeable_dimensions: set, maximum_chunk: int
+):
 
-    # How large is the largest chunk in the input
-    maximum_chunk = reduce(mul, map(max, x.chunks))
-
-    changeable_dimensions = set(range(len(x.chunks))) - {axis}
-    new_chunks = list(x.chunks)
-    new_chunks[axis] = (longest_group,)
+    chunksize_tolerance = config.get("array.chunk-size-tolerance")
 
     # iterate until we distributed the increase in chunksize accross all dimensions
     # or every non-shuffle dimension is all 1
@@ -125,7 +114,7 @@ def _rechunk_other_dimensions(
             up_chunksize_limit_for_dim = max(new_chunks[i]) / (
                 chunksize_inc_factor ** (1 / n_changeable_dimensions)
             )
-            for c in x.chunks[i]:
+            for c in input_chunks[i]:
                 if c > chunksize_tolerance * up_chunksize_limit_for_dim:
                     factor = math.ceil(c / up_chunksize_limit_for_dim)
 
@@ -146,7 +135,29 @@ def _rechunk_other_dimensions(
                 changeable_dimensions.remove(i)
 
             new_chunks[i] = tuple(new_chunksizes)
+    return new_chunks
 
+
+def _rechunk_other_dimensions(
+    x: Array, longest_group: int, axis: int, chunks: Literal["auto"]
+) -> Array:
+    assert chunks == "auto", "Only auto is supported for now"
+    chunksize_tolerance = config.get("array.chunk-size-tolerance")
+
+    if longest_group <= max(x.chunks[axis]) * chunksize_tolerance:
+        # We are staying below our threshold, so don't rechunk
+        return x
+
+    changeable_dimensions = set(range(len(x.chunks))) - {axis}
+    new_chunks = list(x.chunks)
+    new_chunks[axis] = (longest_group,)
+
+    # How large is the largest chunk in the input
+    maximum_chunk = reduce(mul, map(max, x.chunks))
+
+    new_chunks = _calculate_new_chunksizes(
+        x.chunks, new_chunks, changeable_dimensions, maximum_chunk
+    )
     new_chunks[axis] = x.chunks[axis]
     return x.rechunk(tuple(new_chunks))
 
diff --git a/dask/array/core.py b/dask/array/core.py
index faedbb6ef..eb6f40723 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -23,8 +23,9 @@ from typing import Any, Literal, TypeVar, Union, cast
 import numpy as np
 from numpy.typing import ArrayLike
 from packaging.version import Version
-from tlz import accumulate, concat, first, frequencies, groupby, partition
+from tlz import accumulate, concat, first, groupby, partition
 from tlz.curried import pluck
+from toolz import frequencies
 
 from dask import compute, config, core
 from dask.array import chunk
@@ -39,6 +40,7 @@ from dask.array.dispatch import (  # noqa: F401
 )
 from dask.array.numpy_compat import NUMPY_GE_200, _Recurser
 from dask.array.slicing import replace_ellipsis, setitem_array, slice_array
+from dask.array.utils import compute_meta, meta_from_array
 from dask.base import (
     DaskMethodsMixin,
     compute_as_if_collection,
@@ -3139,17 +3141,7 @@ def normalize_chunks(chunks, shape=None, limit=None, dtype=None, previous_chunks
     if chunks and shape is not None:
         # allints: did we start with chunks as a simple tuple of ints?
         allints = all(isinstance(c, int) for c in chunks)
-        chunks = sum(
-            (
-                (
-                    blockdims_from_blockshape((s,), (c,))
-                    if not isinstance(c, (tuple, list))
-                    else (c,)
-                )
-                for s, c in zip(shape, chunks)
-            ),
-            (),
-        )
+        chunks = _convert_int_chunk_to_tuple(shape, chunks)
     for c in chunks:
         if not c:
             raise ValueError(
@@ -3179,6 +3171,20 @@ def normalize_chunks(chunks, shape=None, limit=None, dtype=None, previous_chunks
     )
 
 
+def _convert_int_chunk_to_tuple(shape, chunks):
+    return sum(
+        (
+            (
+                blockdims_from_blockshape((s,), (c,))
+                if not isinstance(c, (tuple, list))
+                else (c,)
+            )
+            for s, c in zip(shape, chunks)
+        ),
+        (),
+    )
+
+
 def _compute_multiplier(limit: int, dtype, largest_block: int, result):
     """
     Utility function for auto_chunk, to fin how much larger or smaller the ideal
@@ -3188,7 +3194,7 @@ def _compute_multiplier(limit: int, dtype, largest_block: int, result):
         limit
         / dtype.itemsize
         / largest_block
-        / math.prod(r for r in result.values() if r)
+        / math.prod(max(r) if isinstance(r, tuple) else r for r in result.values() if r)
     )
 
 
@@ -3217,9 +3223,7 @@ def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
     normalize_chunks: for full docstring and parameters
     """
     if previous_chunks is not None:
-        previous_chunks = tuple(
-            c if isinstance(c, tuple) else (c,) for c in previous_chunks
-        )
+        previous_chunks = _convert_int_chunk_to_tuple(shape, previous_chunks)
     chunks = list(chunks)
 
     autos = {i for i, c in enumerate(chunks) if c == "auto"}
@@ -3253,6 +3257,7 @@ def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
             )
 
     limit = max(1, limit)
+    chunksize_tolerance = config.get("array.chunk-size-tolerance")
 
     largest_block = math.prod(
         cs if isinstance(cs, Number) else max(cs) for cs in chunks if cs != "auto"
@@ -3260,7 +3265,11 @@ def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
 
     if previous_chunks:
         # Base ideal ratio on the median chunk size of the previous chunks
-        result = {a: np.median(previous_chunks[a]) for a in autos}
+        median_chunks = {a: np.median(previous_chunks[a]) for a in autos}
+        result = {}
+
+        # How much larger or smaller the ideal chunk size is relative to what we have now
+        multiplier = _compute_multiplier(limit, dtype, largest_block, median_chunks)
 
         ideal_shape = []
         for i, s in enumerate(shape):
@@ -3271,36 +3280,55 @@ def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
             else:
                 ideal_shape.append(s)
 
-        # How much larger or smaller the ideal chunk size is relative to what we have now
-        multiplier = _compute_multiplier(limit, dtype, largest_block, result)
-
-        last_multiplier = 0
-        last_autos = set()
-        while (
-            multiplier != last_multiplier or autos != last_autos
-        ):  # while things change
-            last_multiplier = multiplier  # record previous values
+        def _trivial_aggregate(a):
+            autos.remove(a)
+            del median_chunks[a]
+            return True
+
+        multiplier_remaining = True
+        while multiplier_remaining:  # while things change
             last_autos = set(autos)  # record previous values
+            multiplier_remaining = False
 
             # Expand or contract each of the dimensions appropriately
             for a in sorted(autos):
-                if ideal_shape[a] == 0:
-                    result[a] = 0
-                    continue
-                proposed = result[a] * multiplier ** (1 / len(autos))
+                this_multiplier = multiplier ** (1 / len(last_autos))
+
+                proposed = median_chunks[a] * this_multiplier
+                this_chunksize_tolerance = chunksize_tolerance ** (1 / len(last_autos))
+                max_chunk_size = proposed * this_chunksize_tolerance
+
                 if proposed > shape[a]:  # we've hit the shape boundary
-                    autos.remove(a)
-                    largest_block *= shape[a]
                     chunks[a] = shape[a]
-                    del result[a]
-                else:
+                    multiplier_remaining = _trivial_aggregate(a)
+                    largest_block *= shape[a]
+                    continue
+                elif this_multiplier <= 1 or max(previous_chunks[a]) > max_chunk_size:
                     result[a] = round_to(proposed, ideal_shape[a])
+                    continue
+                else:
+                    dimension_result, new_chunk = [], 0
+                    for c in previous_chunks[a]:
+                        if c + new_chunk <= proposed:
+                            # keep increasing the chunk
+                            new_chunk += c
+                        else:
+                            # We reach the boundary so start a new chunk
+                            dimension_result.append(new_chunk)
+                            new_chunk = c
+                    if new_chunk > 0:
+                        dimension_result.append(new_chunk)
+
+                result[a] = tuple(dimension_result)
 
             # recompute how much multiplier we have left, repeat
-            multiplier = _compute_multiplier(limit, dtype, largest_block, result)
+            if multiplier_remaining:
+                multiplier = _compute_multiplier(
+                    limit, dtype, largest_block, median_chunks
+                )
 
         for k, v in result.items():
-            chunks[k] = v
+            chunks[k] = v if v else 0
         return tuple(chunks)
 
     else:
@@ -3322,6 +3350,25 @@ def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
         return tuple(chunks)
 
 
+def _split_up_single_chunk(
+    c: int, this_chunksize_tolerance: float, max_chunk_size: int, proposed: int
+) -> list[int]:
+    # Calculate by what factor we have to split this chunk
+    m = c / proposed
+    if math.ceil(m) / m > this_chunksize_tolerance:
+        # We want to smooth things potentially if rounding up would change the result
+        # by a lot
+        m = math.ceil(c / max_chunk_size)
+    else:
+        m = math.ceil(m)
+    # split the chunk
+    new_c, remainder = divmod(c, min(m, c))
+    x = [new_c] * min(m, c)
+    for i in range(remainder):
+        x[i] += 1
+    return x
+
+
 def round_to(c, s):
     """Return a chunk dimension that is close to an even multiple or factor
 
@@ -5953,4 +6000,3 @@ class BlockView:
 
 
 from dask.array.blockwise import blockwise
-from dask.array.utils import compute_meta, meta_from_array
diff --git a/dask/array/einsumfuncs.py b/dask/array/einsumfuncs.py
index 1dc8118a9..99601e075 100644
--- a/dask/array/einsumfuncs.py
+++ b/dask/array/einsumfuncs.py
@@ -1,12 +1,17 @@
 from __future__ import annotations
 
+from functools import reduce
+from operator import mul
+
 import numpy as np
 
+from dask.array._shuffle import _calculate_new_chunksizes
 from dask.array.core import asarray, blockwise, einsum_lookup
 from dask.utils import derived_from
 
 einsum_symbols = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
 einsum_symbols_set = set(einsum_symbols)
+from dask import config
 
 
 def chunk_einsum(*operands, **kwargs):
@@ -229,6 +234,39 @@ def einsum(*operands, dtype=None, optimize=False, split_every=None, **kwargs):
     contract_inds = all_inds - set(outputs)
     ncontract_inds = len(contract_inds)
 
+    if len(inputs) > 1 and len(outputs) > 0:
+        # Calculate the increase in chunk size compared to the largest input chunk
+        max_chunk_sizes, max_chunk_size_inuput = [], 0
+        for op, input in zip(ops, inputs):
+            max_chunk_size_inuput = max(
+                reduce(mul, map(max, op.chunks)), max_chunk_size_inuput
+            )
+            max_chunk_sizes.extend(
+                max(op.chunks[i])
+                for i, inp in enumerate(input)
+                if inp not in contract_inds
+            )
+
+        max_chunk_size_output = reduce(mul, max_chunk_sizes)
+        factor = max_chunk_size_output / (
+            max_chunk_size_inuput * config.get("array.chunk-size-tolerance")
+        )
+
+        # Rechunk inputs to make input chunks smaller to avoid an increase in
+        # output chunks
+        new_ops = []
+        for op, input in zip(ops, inputs):
+            changeable_dimensions = {ctr for ctr, i in enumerate(input) if i in outputs}
+            f = factor ** (len(changeable_dimensions) / len(outputs))
+            result = _calculate_new_chunksizes(
+                op.chunks,
+                list(op.chunks),
+                changeable_dimensions,
+                reduce(mul, map(max, op.chunks)) / f,
+            )
+            new_ops.append(op.rechunk(result))
+        ops = new_ops
+
     # Introduce the contracted indices into the blockwise product
     # so that we get numpy arrays, not lists
     result = blockwise(
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
index c18e52939..ec12479cf 100644
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -195,7 +195,9 @@ def slice_with_newaxes(out_name, in_name, blockdims, index):
             where_none[i] -= n
 
     # Pass down and do work
-    dsk, blockdims2 = slice_wrap_lists(out_name, in_name, blockdims, index2)
+    dsk, blockdims2 = slice_wrap_lists(
+        out_name, in_name, blockdims, index2, not where_none
+    )
 
     if where_none:
         expand = expander(where_none)
@@ -220,7 +222,7 @@ def slice_with_newaxes(out_name, in_name, blockdims, index):
         return dsk, blockdims2
 
 
-def slice_wrap_lists(out_name, in_name, blockdims, index):
+def slice_wrap_lists(out_name, in_name, blockdims, index, allow_getitem_optimization):
     """
     Fancy indexing along blocked array dasks
 
@@ -251,7 +253,9 @@ def slice_wrap_lists(out_name, in_name, blockdims, index):
 
     # No lists, hooray! just use slice_slices_and_integers
     if not where_list:
-        return slice_slices_and_integers(out_name, in_name, blockdims, index)
+        return slice_slices_and_integers(
+            out_name, in_name, blockdims, index, allow_getitem_optimization
+        )
 
     # Replace all lists with full slices  [3, 1, 0] -> slice(None, None, None)
     index_without_list = tuple(
@@ -269,7 +273,11 @@ def slice_wrap_lists(out_name, in_name, blockdims, index):
         # Do first pass without lists
         tmp = "slice-" + tokenize((out_name, in_name, blockdims, index))
         dsk, blockdims2 = slice_slices_and_integers(
-            tmp, in_name, blockdims, index_without_list
+            tmp,
+            in_name,
+            blockdims,
+            index_without_list,
+            allow_getitem_optimization=False,
         )
 
         # After collapsing some axes due to int indices, adjust axis parameter
@@ -285,7 +293,9 @@ def slice_wrap_lists(out_name, in_name, blockdims, index):
     return dsk3, blockdims2
 
 
-def slice_slices_and_integers(out_name, in_name, blockdims, index):
+def slice_slices_and_integers(
+    out_name, in_name, blockdims, index, allow_getitem_optimization=False
+):
     """
     Dask array indexing with slices and integers
 
@@ -329,7 +339,12 @@ def slice_slices_and_integers(out_name, in_name, blockdims, index):
     all_slices = list(product(*[pluck(1, s) for s in sorted_block_slices]))
 
     dsk_out = {
-        out_name: (getitem, in_name, slices)
+        out_name: (
+            (getitem, in_name, slices)
+            if not allow_getitem_optimization
+            or not all(sl == slice(None, None, None) for sl in slices)
+            else in_name
+        )
         for out_name, in_name, slices in zip(out_names, in_names, all_slices)
     }
 
diff --git a/dask/array/utils.py b/dask/array/utils.py
index 1c7adc325..64a98bc2e 100644
--- a/dask/array/utils.py
+++ b/dask/array/utils.py
@@ -10,7 +10,6 @@ import warnings
 import numpy as np
 from tlz import concat, frequencies
 
-from dask.array.core import Array
 from dask.array.numpy_compat import AxisError
 from dask.base import is_dask_collection, tokenize
 from dask.highlevelgraph import HighLevelGraph
@@ -440,6 +439,8 @@ def arange_safe(*args, like, **kwargs):
 
 
 def _array_like_safe(np_func, da_func, a, like, **kwargs):
+    from dask.array.core import Array
+
     if like is a and hasattr(a, "__array_function__"):
         return a
 
diff --git a/dask/order.py b/dask/order.py
index d80220338..5f678833a 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -562,7 +562,7 @@ def order(
         while next_deps:
             item = max(next_deps, key=sort_key)
             path_append(item)
-            next_deps = dependencies[item]
+            next_deps = dependencies[item].difference(result)
             path_extend(next_deps)
 
         # B. Walk the critical path
@@ -670,12 +670,12 @@ def _connecting_to_roots(
                 and (result_first is r_child or r_child.issubset(result_first))
             ):
                 identical_sets = False
-                if not new_set:
+                if new_set is None:
                     new_set = set(result_first)
                 max_dependents[key] = max(max_dependents[child], max_dependents[key])
                 new_set.update(r_child)
 
-        if new_set:
+        if new_set is not None:
             new_set_frozen = frozenset(new_set)
             deduped = dedup_mapping.get(new_set_frozen, None)
             if deduped is None:


Test patch:
diff --git a/dask/array/tests/test_rechunk.py b/dask/array/tests/test_rechunk.py
index 02eea953f..3ac4d71bb 100644
--- a/dask/array/tests/test_rechunk.py
+++ b/dask/array/tests/test_rechunk.py
@@ -825,7 +825,19 @@ def test_rechunk_avoid_needless_chunking():
         (100, 50, 10, (10,) * 10),
         (100, 100, 10, (10,) * 10),
         (20, 7, 10, (7, 7, 6)),
-        (20, (1, 1, 1, 1, 6, 2, 1, 7), 5, (5, 5, 5, 5)),
+        (
+            20,
+            (1, 1, 1, 1, 6, 2, 1, 7),
+            5,
+            (5, 5, 5, 5),
+        ),  # This should be smarter
+        (21, (1,) * 21, 10, (10, 10, 1)),  # ensure that we squash together properly
+        (30, (3,) * 10, 7, (6,) * 5),  # ensure that we squash together properly
+        (20, ((2, 2, 2, 9, 1, 2, 2)), 5, (4, 4, 4, 4, 4)),
+        (20, ((10, 10)), 4, (4,) * 5),  # this should split branches evenly
+        (21, ((10, 11)), 4, (4, 4, 4, 4, 4, 1)),  # this should split branches evenly
+        (20, ((1, 18, 1)), 5, (5,) * 4),
+        (38, ((10, 18, 10)), 5, (5, 5, 5, 5, 5, 5, 5, 3)),
     ],
 )
 def test_rechunk_auto_1d(shape, chunks, bs, expected):
@@ -834,6 +846,26 @@ def test_rechunk_auto_1d(shape, chunks, bs, expected):
     assert y.chunks == (expected,)
 
 
+@pytest.mark.parametrize(
+    "previous_chunks,bs,expected",
+    [
+        (((1, 1, 1), (10, 10, 10, 10, 10, 10, 10, 10)), 160, ((3,), (50, 30))),
+        (((2, 2), (20,)), 5, ((1, 1, 1, 1), (7, 7, 6))),
+        (((1, 1), (20,)), 5, ((1, 1), (10, 10))),
+    ],
+)
+def test_normalize_chunks_auto_2d(previous_chunks, bs, expected):
+    shape = tuple(map(sum, previous_chunks))
+    result = normalize_chunks(
+        {0: "auto", 1: "auto"},
+        shape,
+        limit=bs,
+        dtype=np.int8,
+        previous_chunks=previous_chunks,
+    )
+    assert result == expected
+
+
 def test_rechunk_auto_2d():
     x = da.ones((20, 20), chunks=(2, 2))
     y = x.rechunk({0: -1, 1: "auto"}, block_size_limit=20 * x.dtype.itemsize)
@@ -880,7 +912,12 @@ def test_rechunk_auto_image_stack(n):
     with dask.config.set({"array.chunk-size": "1MiB"}):
         x = da.ones((n, 1000, 1000), chunks=(1, 1000, 1000), dtype="float64")
         z = x.rechunk("auto")
-        assert z.chunks == ((1,) * n, (362, 362, 276), (362, 362, 276))
+        assert z.chunks == ((1,) * n, (507, 493), (507, 493))
+
+    with dask.config.set({"array.chunk-size": "1MiB"}):
+        x = da.ones((n, 2000, 2000), chunks=(1, 1000, 1000), dtype="float64")
+        z = x.rechunk("auto")
+        assert z.chunks == ((1,) * n, (507, 507, 507, 479), (507, 507, 507, 479))
 
 
 def test_rechunk_down():
@@ -891,7 +928,7 @@ def test_rechunk_down():
 
     with dask.config.set({"array.chunk-size": "1MiB"}):
         z = y.rechunk("auto")
-        assert z.chunks == ((4,) * 25, (511, 489), (511, 489))
+        assert z.chunks == ((4,) * 25, (471, 471, 58), (471, 471, 58))
 
     with dask.config.set({"array.chunk-size": "1MiB"}):
         z = y.rechunk({0: "auto"})
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
index 9223e356f..6871913f5 100644
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -2455,6 +2455,42 @@ def test_einsum(einsum_signature):
         )
 
 
+def test_einsum_chunksizes():
+    arr1 = da.random.random((1024, 16, 16, 16, 16), chunks=(256, 8, 8, 8, 8))
+    arr2 = da.random.random((1024, 16, 16, 16, 16), chunks=(256, 8, 8, 8, 8))
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=da.PerformanceWarning)
+        result = da.einsum("aijkl,amnop->ijklmnop", arr1, arr2)
+    assert result.chunks == ((4,) * 4,) * 8
+
+    arr1 = da.random.random((64, 16, 16, 16, 16), chunks=(32, 8, 1, 8, 8))
+    arr2 = da.random.random((64, 16, 16, 16, 16), chunks=(32, 8, 8, 1, 8))
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=da.PerformanceWarning)
+        result = da.einsum("aijkl,amnop->ijklmnop", arr1, arr2)
+    assert result.chunks == (
+        (4,) * 4,
+        (1,) * 16,
+        (4,) * 4,
+        (4,) * 4,
+        (4,) * 4,
+        (4,) * 4,
+        (1,) * 16,
+        (4,) * 4,
+    )
+
+    np_arr1 = np.random.random((2, 4, 4, 4))
+    np_arr2 = np.random.random((2, 4, 4, 4))
+
+    arr1 = da.from_array(np_arr1, chunks=(1, 2, 2, 2))
+    arr2 = da.from_array(np_arr2, chunks=(1, 2, 2, 2))
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=da.PerformanceWarning)
+        result = da.einsum("aijk,amno->ijkmno", arr1, arr2)
+    assert result.chunks == ((1,) * 4,) * 6
+    assert_eq(np.einsum("aijk,amno->ijkmno", np_arr1, np_arr2), result)
+
+
 @pytest.mark.parametrize(
     "optimize_opts", [(True, False), ("greedy", False), ("optimal", False)]
 )
diff --git a/dask/array/tests/test_slicing.py b/dask/array/tests/test_slicing.py
index a5ea27370..2d1ab3006 100644
--- a/dask/array/tests/test_slicing.py
+++ b/dask/array/tests/test_slicing.py
@@ -1036,3 +1036,13 @@ def test_take_sorted_indexer():
             )
         },
     } == dict(result.dask)
+
+
+def test_all_none_slices_just_mappings():
+    arr = da.ones((10, 10), chunks=(1, 5))
+    result = arr[slice(None, 6), slice(None)]
+    dsk = dict(result.dask)
+    assert len([k for k in dsk if "getitem" in k[0]]) == 12
+    # check that we are just mapping the keys
+    assert all(v in dsk for k, v in dsk.items() if "getitem" in k[0])
+    assert_eq(result, np.ones((6, 10)))
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
index f28e21f13..2078e494b 100644
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -546,8 +546,6 @@ def test_blockwise_array_creation(c, io, fuse):
 def test_blockwise_dataframe_io(c, tmpdir, io, fuse, from_futures):
     pd = pytest.importorskip("pandas")
     dd = pytest.importorskip("dask.dataframe")
-    if dd._dask_expr_enabled():
-        pytest.xfail("doesn't work yet")
 
     df = pd.DataFrame({"x": [1, 2, 3] * 5, "y": range(15)})
 
diff --git a/dask/tests/test_optimization.py b/dask/tests/test_optimization.py
index 800a17f0c..eca7fe4d2 100644
--- a/dask/tests/test_optimization.py
+++ b/dask/tests/test_optimization.py
@@ -1542,7 +1542,7 @@ def test_fusion_legacy_hybrid():
 
     new_dsk = convert_legacy_graph(fused)
     assert isinstance(new_dsk["bar"], Alias)
-    assert new_dsk["bar"].key == fused_task_key
+    assert new_dsk["bar"].key == "bar"
     assert not new_dsk[fused_task_key].dependencies
     assert new_dsk[fused_task_key]() == "a-b=c"
 
@@ -1590,7 +1590,7 @@ def test_fusion_wide_legacy_hybrid():
 
         new_dsk = convert_legacy_graph(fused)
         assert isinstance(new_dsk["bar"], Alias)
-        assert new_dsk["bar"].key == fused_task_key
+        assert new_dsk["bar"].key == "bar"
         assert not new_dsk[fused_task_key].dependencies
         assert new_dsk[fused_task_key]() == "a-b=a-b"
 
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index 0604d95aa..51f83e691 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -84,6 +84,32 @@ def test_repr():
     assert repr(t) == "<Task 'kwarg' use_kwargs('foo', kwarg='kwarg_value')>"
 
 
+def test_unpickled_repr():
+    t = pickle.loads(pickle.dumps(Task("key", func, "a", "b")))
+    assert repr(t) == "<Task 'key' func('a', 'b')>"
+
+    t = pickle.loads(pickle.dumps(Task("nested", func2, t, t.ref())))
+    assert (
+        repr(t) == "<Task 'nested' func2(<Task 'key' func('a', 'b')>, Alias(key->key))>"
+    )
+
+    def long_function_name_longer_even_longer(a, b):
+        return a + b
+
+    t = pickle.loads(
+        pickle.dumps(Task("long", long_function_name_longer_even_longer, t, t.ref()))
+    )
+    assert repr(t) == "<Task 'long' long_function_name_longer_even_longer(...)>"
+
+    def use_kwargs(a, kwarg=None):
+        return a + kwarg
+
+    t = pickle.loads(
+        pickle.dumps(Task("kwarg", use_kwargs, "foo", kwarg="kwarg_value"))
+    )
+    assert repr(t) == "<Task 'kwarg' use_kwargs('foo', kwarg='kwarg_value')>"
+
+
 def _assert_dsk_conversion(new_dsk):
     vals = [(k, v.key) for k, v in new_dsk.items()]
     assert all(v[0] == v[1] for v in vals), vals
@@ -635,6 +661,51 @@ def test_resolve_aliases():
     assert optimized == dsk
 
 
+def test_resolve_multiple_aliases():
+
+    tasks = [
+        Task("first", func, 10),
+        Alias("second", "first"),
+        Task("third", func, TaskRef("second")),
+        Alias("fourth", "third"),
+        Task("fifth", func, TaskRef("fourth")),
+    ]
+    dsk = {t.key: t for t in tasks}
+    assert len(dsk) == 5
+
+    optimized = resolve_aliases(dsk, {"fifth"}, reverse_dict(DependenciesMapping(dsk)))
+    assert len(optimized) == 3
+    expected = dsk["third"].copy()
+    expected.key = "fourth"
+    assert optimized["fourth"] == expected
+
+    expected = dsk["first"].copy()
+    expected.key = "second"
+    assert optimized["second"] == expected
+
+
+def test_convert_resolve():
+    dsk = {
+        "first": (func, 10),
+        "second": "first",
+        "third": (func, "second"),
+        "fourth": "third",
+        "fifth": (func, "fourth"),
+    }
+    dsk = convert_legacy_graph(dsk)
+    assert len(dsk) == 5
+
+    optimized = resolve_aliases(dsk, {"fifth"}, reverse_dict(DependenciesMapping(dsk)))
+    assert len(optimized) == 3
+    expected = dsk["third"].copy()
+    expected.key = "fourth"
+    assert optimized["fourth"] == expected
+
+    expected = dsk["first"].copy()
+    expected.key = "second"
+    assert optimized["second"] == expected
+
+
 def test_parse_nested():
     t = Task(
         "key",