Repository: dask/dask
Pull request #11687 (2025-01-21 13:55:20)

Reduce pickle size for tasks
This makes the tasks a bit leaner by reducing pickle size. The performance gain is rather secondary but this should save a lot of memory client side when shipping things to the scheduler. 

(I doubt that it will save anything on the scheduler since the pool for the pickle byte stream will be eaten up by order and the scheduler state but I haven't tried yet)

Two things left
- ~I would like to avoid that every subclass of Task has to overload the getstate/setstate which is why I added the BaseTask. Child classes don't necessarily need this complexity. For example, the shuffle barrier task we're using as a subclass is created once per shuffle and a couple bytes overhead is not that bad. The option I chose here is very ugly. I'm still exploring options. Open for suggestions (This is why linting is still broken)~
- ~A part of me wonders if we can write a pseudo PickleBuffer object that serializes this stuff on demand. That would help a lot while we're still materializing the graph fully client side. I will explore this briefly... I suspect this is possible but I expect it to break memoization. If somebody already knows this is stupid, please stop me!~ Answer: Tricky. Couldn't find a way to do it. the picklebuffer interface is really strict (but I haven't spent much time ont his. We'll hopefully avoid client side materialization soon)

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

     15 files  ± 0       15 suites  ±0   4h 27m 19s ⏱️ + 1m 59s
 17 142 tests + 1   15 853 ✅  - 108   1 182 💤 +2  107 ❌ +107 
211 227 runs  +15  193 353 ✅  - 838  17 025 💤 +4  849 ❌ +849 

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

Results for commit 16c631d1. ± Comparison against base commit f5bd5590.

<details>
  <summary>This pull request <b>skips</b> 2 tests.</summary>

```
dask.tests.test_distributed ‑ test_to_hdf_scheduler_distributed[10]
dask.tests.test_distributed ‑ test_to_hdf_scheduler_distributed[4]
```
</details>

[test-results]:data:application/gzip;base64,H4sIACmzj2cC/03MSw7CIBSF4a00jB1wgcvDzZgWSkLsw1AYGfcu1NI6PN+F/018mMaN3DvAW0e2HNK1XI59CutSt6TcFCrXtN8VCNb2Y8vW7p808guf4VUR9PXQ92GqRtVJY4xrLEaLxLzUOANgTB37jBvO9/oPj7iiDJsdcS1Mk/+2Xec5pDIISCs5OBB8YBo9NWLgQlrhAd2oQCM6hlYr8vkCWTIdyh4BAAA=

fjetter: Some preliminary tests on small-ish scale using the coiled geobenchmark climatology / highlevel API show that this change reduces the client side memory footprint by about 30% reduction of RSS memory.

The pickled graph size reduces from roughly 330MiB to 130MiB, i.e. that's a 60% reduction. That's still a bit large considering the size (small size in the benchmark suite) but most of this is driven by the graph being materialized fully. Dropping the client side materialization should shrink this dramatically.
fjetter: One test is failing `dask/dataframe/dask_expr/tests/test_datasets.py::test_timeseries_gaph_size[None]`

This is a test that compares graph sizes between legacy and dask-expr timeseries implementations. The dask-expr is typically strictly smaller. Apparently on this windows test that is not true but I suspect this is somehow due to the random seed not being set.

I suspect we'll kill the legacy version of anyhow. I'll retrigger the test job and will remove the legacy version
fjetter: Interestingly, in my local testing, the dask-expr graph is about 15x slimmer than the legacy one. No idea why. Probably some lambdas/local contamination somewhere but I haven't checked
fjetter: Interestingly, the windows thing is persistent. I'll skip that test on this platform for now. I suspect we're going to rip out the legacy timeseries anyhow
fjetter: going to merge once ci is through. @hendrikmakait do you still want to review?

Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 5506e5553..45e6c4667 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -58,7 +58,7 @@ import itertools
 import sys
 from collections import defaultdict
 from collections.abc import Callable, Container, Iterable, Mapping, MutableMapping
-from functools import partial
+from functools import lru_cache, partial
 from typing import Any, TypeVar, cast
 
 from dask.sizeof import sizeof
@@ -402,9 +402,10 @@ class GraphNode:
         return False
 
     def __sizeof__(self) -> int:
-        return sum(
-            sizeof(getattr(self, sl)) for sl in type(self).__slots__
-        ) + sys.getsizeof(type(self))
+        all_slots = self.get_all_slots()
+        return sum(sizeof(getattr(self, sl)) for sl in all_slots) + sys.getsizeof(
+            type(self)
+        )
 
     def substitute(
         self, subs: dict[KeyType, KeyType | GraphNode], key: KeyType | None = None
@@ -475,6 +476,16 @@ class GraphNode:
             _data_producer=any(t.data_producer for t in tasks),
         )
 
+    @classmethod
+    @lru_cache
+    def get_all_slots(cls):
+        slots = list()
+        for c in cls.mro():
+            slots.extend(getattr(c, "__slots__", ()))
+        # Interestingly, sorting this causes the nested containers to pickle
+        # more efficiently
+        return sorted(set(slots))
+
 
 _no_deps: frozenset = frozenset()
 
@@ -498,6 +509,9 @@ class Alias(GraphNode):
         self.target = target
         self._dependencies = frozenset((self.target,))
 
+    def __reduce__(self):
+        return Alias, (self.key, self.target)
+
     def copy(self):
         return Alias(self.key, self.target)
 
@@ -740,6 +754,15 @@ class Task(GraphNode):
             return self.func(*new_argspec, **kwargs)
         return self.func(*new_argspec)
 
+    def __setstate__(self, state):
+        slots = self.__class__.get_all_slots()
+        for sl, val in zip(slots, state):
+            setattr(self, sl, val)
+
+    def __getstate__(self):
+        slots = self.__class__.get_all_slots()
+        return tuple(getattr(self, sl) for sl in slots)
+
     @property
     def is_coro(self):
         if self._is_coro is None:
@@ -810,6 +833,23 @@ class NestedContainer(Task):
             **kwargs,
         )
 
+    def __getstate__(self):
+        state = super().__getstate__()
+        state = list(state)
+        slots = self.__class__.get_all_slots()
+        ix = slots.index("kwargs")
+        # The constructor as a kwarg is redundant since this is encoded in the
+        # class itself. Serializing the builtin types is not trivial
+        # This saves about 15% of overhead
+        state[ix] = state[ix].copy()
+        state[ix].pop("constructor", None)
+        return state
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        self.kwargs["constructor"] = self.__class__.constructor
+        return self
+
     def __repr__(self):
         return f"{type(self).__name__}({self.args})"
 
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
index 1a32a8994..e5447597e 100644
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -349,6 +349,27 @@ if _array_expr_enabled():
             blockwise,
             elemwise,
             random,
+            rechunk,
+            reduction,
+        )
+        from dask.array.reductions import (
+            all,
+            any,
+            max,
+            mean,
+            min,
+            moment,
+            nanmax,
+            nanmean,
+            nanmin,
+            nanprod,
+            nanstd,
+            nansum,
+            nanvar,
+            prod,
+            std,
+            sum,
+            var,
         )
 
         arange = raise_not_implemented_error("arange")
@@ -356,17 +377,7 @@ if _array_expr_enabled():
         empty = raise_not_implemented_error("empty")
         from_array = raise_not_implemented_error("from_array")
         linspace = raise_not_implemented_error("linspace")
-        mean = raise_not_implemented_error("mean")
-        moment = raise_not_implemented_error("moment")
-        nanmean = raise_not_implemented_error("nanmean")
-        nanstd = raise_not_implemented_error("nanstd")
-        nansum = raise_not_implemented_error("nansum")
-        nanvar = raise_not_implemented_error("nanvar")
         ones = raise_not_implemented_error("ones")
-        prod = raise_not_implemented_error("prod")
-        std = raise_not_implemented_error("std")
-        sum = raise_not_implemented_error("sum")
-        var = raise_not_implemented_error("var")
         zeros = raise_not_implemented_error("zeros")
         backends = raise_not_implemented_error("backends")
         fft = raise_not_implemented_error("fft")
@@ -375,7 +386,6 @@ if _array_expr_enabled():
         ma = raise_not_implemented_error("ma")
         overlap = raise_not_implemented_error("overlap")
         atop = raise_not_implemented_error("atop")
-        blockwise = raise_not_implemented_error("blockwise")
         register_chunk_type = raise_not_implemented_error("register_chunk_type")
         asanyarray = raise_not_implemented_error("asanyarray")
         block = raise_not_implemented_error("block")
@@ -414,26 +424,17 @@ if _array_expr_enabled():
         optimize = raise_not_implemented_error("optimize")
         map_overlap = raise_not_implemented_error("map_overlap")
         percentile = raise_not_implemented_error("percentile")
-        rechunk = raise_not_implemented_error("rechunk")
-        all = raise_not_implemented_error("all")
-        any = raise_not_implemented_error("any")
         argmax = raise_not_implemented_error("argmax")
         argmin = raise_not_implemented_error("argmin")
         argtopk = raise_not_implemented_error("argtopk")
         cumprod = raise_not_implemented_error("cumprod")
         cumsum = raise_not_implemented_error("cumsum")
-        max = raise_not_implemented_error("max")
         median = raise_not_implemented_error("median")
-        min = raise_not_implemented_error("min")
         nanargmax = raise_not_implemented_error("nanargmax")
         nanargmin = raise_not_implemented_error("nanargmin")
         nancumprod = raise_not_implemented_error("nancumprod")
         nancumsum = raise_not_implemented_error("nancumsum")
-        nanmax = raise_not_implemented_error("nanmax")
         nanmedian = raise_not_implemented_error("nanmedian")
-        nanmin = raise_not_implemented_error("nanmin")
-        nanprod = raise_not_implemented_error("nanprod")
-        reduction = raise_not_implemented_error("reduction")
         topk = raise_not_implemented_error("topk")
         trace = raise_not_implemented_error("trace")
         reshape = raise_not_implemented_error("reshape")
diff --git a/dask/array/_array_expr/__init__.py b/dask/array/_array_expr/__init__.py
index 745aae6d1..033f6a7ca 100644
--- a/dask/array/_array_expr/__init__.py
+++ b/dask/array/_array_expr/__init__.py
@@ -2,4 +2,5 @@ from __future__ import annotations
 
 import dask.array._array_expr._backends
 from dask.array._array_expr import random
-from dask.array._array_expr._collection import Array, blockwise, elemwise
+from dask.array._array_expr._collection import Array, blockwise, elemwise, rechunk
+from dask.array._array_expr._reductions import _tree_reduce, reduction
diff --git a/dask/array/_array_expr/_blockwise.py b/dask/array/_array_expr/_blockwise.py
index ea9ee7d1d..33e8b027c 100644
--- a/dask/array/_array_expr/_blockwise.py
+++ b/dask/array/_array_expr/_blockwise.py
@@ -7,7 +7,7 @@ import numpy as np
 import tlz as toolz
 
 from dask import is_dask_collection
-from dask.array._array_expr._expr import ArrayExpr
+from dask.array._array_expr._expr import ArrayExpr, unify_chunks_expr
 from dask.array._array_expr._utils import compute_meta
 from dask.array.core import (
     _elemwise_handle_where,
@@ -69,16 +69,21 @@ class Blockwise(ArrayExpr):
 
     @cached_property
     def chunks(self):
-        arginds = [(a, i) for (a, i) in toolz.partition(2, self.args) if i is not None]
-        chunkss = {}
-        # For each dimension, use the input chunking that has the most blocks;
-        # this will ensure that broadcasting works as expected, and in
-        # particular the number of blocks should be correct if the inputs are
-        # consistent.
-        for arg, ind in arginds:
-            for c, i in zip(arg.chunks, ind):
-                if i not in chunkss or len(c) > len(chunkss[i]):
-                    chunkss[i] = c
+        if self.align_arrays:
+            chunkss, arrays, _ = unify_chunks_expr(*self.args)
+        else:
+            arginds = [
+                (a, i) for (a, i) in toolz.partition(2, self.args) if i is not None
+            ]
+            chunkss = {}
+            # For each dimension, use the input chunking that has the most blocks;
+            # this will ensure that broadcasting works as expected, and in
+            # particular the number of blocks should be correct if the inputs are
+            # consistent.
+            for arg, ind in arginds:
+                for c, i in zip(arg.chunks, ind):
+                    if i not in chunkss or len(c) > len(chunkss[i]):
+                        chunkss[i] = c
 
         for k, v in self.new_axes.items():
             if not isinstance(v, tuple):
@@ -178,6 +183,12 @@ class Blockwise(ArrayExpr):
         )
         return dict(graph)
 
+    def _lower(self):
+        if self.align_arrays:
+            _, arrays, changed = unify_chunks_expr(*self.args)
+            if changed:
+                return type(self)(*self.operands[: len(self._parameters)], *arrays)
+
 
 class Elemwise(Blockwise):
     _parameters = ["op", "dtype", "name", "where"]
diff --git a/dask/array/_array_expr/_collection.py b/dask/array/_array_expr/_collection.py
index 4f262cbb5..06ac16606 100644
--- a/dask/array/_array_expr/_collection.py
+++ b/dask/array/_array_expr/_collection.py
@@ -9,7 +9,7 @@ from dask._collections import new_collection
 from dask.array._array_expr._blockwise import Blockwise, Elemwise
 from dask.array._array_expr._expr import ArrayExpr
 from dask.array._array_expr._io import FromGraph
-from dask.array.core import finalize
+from dask.array.core import T_IntOrNaN, finalize
 from dask.base import DaskMethodsMixin, named_schedulers
 from dask.core import flatten
 from dask.utils import key_split
@@ -84,6 +84,14 @@ class Array(DaskMethodsMixin):
     def ndim(self):
         return self.expr.ndim
 
+    @property
+    def numblocks(self):
+        return self.expr.numblocks
+
+    @property
+    def size(self) -> T_IntOrNaN:
+        return self.expr.size
+
     @property
     def name(self):
         return self.expr.name
@@ -128,6 +136,206 @@ class Array(DaskMethodsMixin):
         # TODO(expr-soon): Not done yet, but needed for assert_eq to identify us as an Array
         raise NotImplementedError
 
+    def rechunk(
+        self,
+        chunks="auto",
+        threshold=None,
+        block_size_limit=None,
+        balance=False,
+        method=None,
+    ):
+        return rechunk(self, chunks, threshold, block_size_limit, balance, method)
+
+    def sum(self, axis=None, dtype=None, keepdims=False, split_every=None, out=None):
+        """
+        Return the sum of the array elements over the given axis.
+
+        Refer to :func:`dask.array.sum` for full documentation.
+
+        See Also
+        --------
+        dask.array.sum : equivalent function
+        """
+        from dask.array.reductions import sum
+
+        return sum(
+            self,
+            axis=axis,
+            dtype=dtype,
+            keepdims=keepdims,
+            split_every=split_every,
+            out=out,
+        )
+
+    def mean(self, axis=None, dtype=None, keepdims=False, split_every=None, out=None):
+        """Returns the average of the array elements along given axis.
+
+        Refer to :func:`dask.array.mean` for full documentation.
+
+        See Also
+        --------
+        dask.array.mean : equivalent function
+        """
+        from dask.array.reductions import mean
+
+        return mean(
+            self,
+            axis=axis,
+            dtype=dtype,
+            keepdims=keepdims,
+            split_every=split_every,
+            out=out,
+        )
+
+    def std(
+        self, axis=None, dtype=None, keepdims=False, ddof=0, split_every=None, out=None
+    ):
+        """Returns the standard deviation of the array elements along given axis.
+
+        Refer to :func:`dask.array.std` for full documentation.
+
+        See Also
+        --------
+        dask.array.std : equivalent function
+        """
+        from dask.array.reductions import std
+
+        return std(
+            self,
+            axis=axis,
+            dtype=dtype,
+            keepdims=keepdims,
+            ddof=ddof,
+            split_every=split_every,
+            out=out,
+        )
+
+    def var(
+        self, axis=None, dtype=None, keepdims=False, ddof=0, split_every=None, out=None
+    ):
+        """Returns the variance of the array elements, along given axis.
+
+        Refer to :func:`dask.array.var` for full documentation.
+
+        See Also
+        --------
+        dask.array.var : equivalent function
+        """
+        from dask.array.reductions import var
+
+        return var(
+            self,
+            axis=axis,
+            dtype=dtype,
+            keepdims=keepdims,
+            ddof=ddof,
+            split_every=split_every,
+            out=out,
+        )
+
+    def moment(
+        self,
+        order,
+        axis=None,
+        dtype=None,
+        keepdims=False,
+        ddof=0,
+        split_every=None,
+        out=None,
+    ):
+        """Calculate the nth centralized moment.
+
+        Refer to :func:`dask.array.moment` for the full documentation.
+
+        See Also
+        --------
+        dask.array.moment : equivalent function
+        """
+        from dask.array.reductions import moment
+
+        return moment(
+            self,
+            order,
+            axis=axis,
+            dtype=dtype,
+            keepdims=keepdims,
+            ddof=ddof,
+            split_every=split_every,
+            out=out,
+        )
+
+    def prod(self, axis=None, dtype=None, keepdims=False, split_every=None, out=None):
+        """Return the product of the array elements over the given axis
+
+        Refer to :func:`dask.array.prod` for full documentation.
+
+        See Also
+        --------
+        dask.array.prod : equivalent function
+        """
+        from dask.array.reductions import prod
+
+        return prod(
+            self,
+            axis=axis,
+            dtype=dtype,
+            keepdims=keepdims,
+            split_every=split_every,
+            out=out,
+        )
+
+    def any(self, axis=None, keepdims=False, split_every=None, out=None):
+        """Returns True if any of the elements evaluate to True.
+
+        Refer to :func:`dask.array.any` for full documentation.
+
+        See Also
+        --------
+        dask.array.any : equivalent function
+        """
+        from dask.array.reductions import any
+
+        return any(self, axis=axis, keepdims=keepdims, split_every=split_every, out=out)
+
+    def all(self, axis=None, keepdims=False, split_every=None, out=None):
+        """Returns True if all elements evaluate to True.
+
+        Refer to :func:`dask.array.all` for full documentation.
+
+        See Also
+        --------
+        dask.array.all : equivalent function
+        """
+        from dask.array.reductions import all
+
+        return all(self, axis=axis, keepdims=keepdims, split_every=split_every, out=out)
+
+    def min(self, axis=None, keepdims=False, split_every=None, out=None):
+        """Return the minimum along a given axis.
+
+        Refer to :func:`dask.array.min` for full documentation.
+
+        See Also
+        --------
+        dask.array.min : equivalent function
+        """
+        from dask.array.reductions import min
+
+        return min(self, axis=axis, keepdims=keepdims, split_every=split_every, out=out)
+
+    def max(self, axis=None, keepdims=False, split_every=None, out=None):
+        """Return the maximum along a given axis.
+
+        Refer to :func:`dask.array.max` for full documentation.
+
+        See Also
+        --------
+        dask.array.max : equivalent function
+        """
+        from dask.array.reductions import max
+
+        return max(self, axis=axis, keepdims=keepdims, split_every=split_every, out=out)
+
 
 def from_graph(layer, _meta, chunks, keys, name_prefix):
     return new_collection(
@@ -150,7 +358,7 @@ def blockwise(
     dtype=None,
     adjust_chunks=None,
     new_axes=None,
-    align_arrays=False,  # TODO(expr-soon): Change to True
+    align_arrays=True,
     concatenate=None,
     meta=None,
     **kwargs,
@@ -291,7 +499,6 @@ def blockwise(
     array([[1235, 1236],
            [1237, 1238]])
     """
-    assert not align_arrays  # TODO(expr-soon): Remove
     new_axes = new_axes or {}
 
     # Input Validation
@@ -376,3 +583,75 @@ def elemwise(op, *args, out=None, where=True, dtype=None, name=None, **kwargs):
     args = [np.asarray(a) if isinstance(a, (list, tuple)) else a for a in args]
 
     return new_collection(Elemwise(op, dtype, name, where, *args))
+
+
+def rechunk(
+    x,
+    chunks="auto",
+    threshold=None,
+    block_size_limit=None,
+    balance=False,
+    method=None,
+):
+    """
+    Convert blocks in dask array x for new chunks.
+
+    Parameters
+    ----------
+    x: dask array
+        Array to be rechunked.
+    chunks:  int, tuple, dict or str, optional
+        The new block dimensions to create. -1 indicates the full size of the
+        corresponding dimension. Default is "auto" which automatically
+        determines chunk sizes.
+    threshold: int, optional
+        The graph growth factor under which we don't bother introducing an
+        intermediate step.
+    block_size_limit: int, optional
+        The maximum block size (in bytes) we want to produce
+        Defaults to the configuration value ``array.chunk-size``
+    balance : bool, default False
+        If True, try to make each chunk to be the same size.
+
+        This means ``balance=True`` will remove any small leftover chunks, so
+        using ``x.rechunk(chunks=len(x) // N, balance=True)``
+        will almost certainly result in ``N`` chunks.
+    method: {'tasks', 'p2p'}, optional.
+        Rechunking method to use.
+
+
+    Examples
+    --------
+    >>> import dask.array as da
+    >>> x = da.ones((1000, 1000), chunks=(100, 100))
+
+    Specify uniform chunk sizes with a tuple
+
+    >>> y = x.rechunk((1000, 10))
+
+    Or chunk only specific dimensions with a dictionary
+
+    >>> y = x.rechunk({0: 1000})
+
+    Use the value ``-1`` to specify that you want a single chunk along a
+    dimension or the value ``"auto"`` to specify that dask can freely rechunk a
+    dimension to attain blocks of a uniform block size
+
+    >>> y = x.rechunk({0: -1, 1: 'auto'}, block_size_limit=1e8)
+
+    If a chunk size does not divide the dimension then rechunk will leave any
+    unevenness to the last chunk.
+
+    >>> x.rechunk(chunks=(400, -1)).chunks
+    ((400, 400, 200), (1000,))
+
+    However if you want more balanced chunks, and don't mind Dask choosing a
+    different chunksize for you then you can use the ``balance=True`` option.
+
+    >>> x.rechunk(chunks=(400, -1), balance=True).chunks
+    ((500, 500), (1000,))
+    """
+
+    return new_collection(
+        x.expr.rechunk(chunks, threshold, block_size_limit, balance, method)
+    )
diff --git a/dask/array/_array_expr/_expr.py b/dask/array/_array_expr/_expr.py
index 62acd4375..017dca926 100644
--- a/dask/array/_array_expr/_expr.py
+++ b/dask/array/_array_expr/_expr.py
@@ -1,9 +1,14 @@
 from __future__ import annotations
 
-from functools import cached_property
+from functools import cached_property, reduce
+from operator import mul
+
+import numpy as np
+import toolz
 
 from dask._expr import Expr
-from dask.array.core import T_IntOrNaN
+from dask.array.core import T_IntOrNaN, common_blockdim
+from dask.blockwise import broadcast_dimensions
 from dask.utils import cached_cumsum
 
 
@@ -38,6 +43,11 @@ class ArrayExpr(Expr):
     def numblocks(self):
         return tuple(map(len, self.chunks))
 
+    @cached_property
+    def size(self) -> T_IntOrNaN:
+        """Number of elements in array"""
+        return reduce(mul, self.shape, 1)
+
     @property
     def name(self):
         return self._name
@@ -93,3 +103,65 @@ class ArrayExpr(Expr):
                 "This often means that you are attempting to use an unsupported "
                 f"API function.."
             )
+
+    def rechunk(
+        self,
+        chunks="auto",
+        threshold=None,
+        block_size_limit=None,
+        balance=False,
+        method=None,
+    ):
+        if self.ndim > 0 and all(s == 0 for s in self.shape):
+            return self
+
+        from dask.array._array_expr._rechunk import Rechunk
+
+        return Rechunk(self, chunks, threshold, block_size_limit, balance, method)
+
+
+def unify_chunks_expr(*args):
+    # TODO(expr): This should probably be a dedicated expression
+    # This is the implementation that expects the inputs to be expressions, the public facing
+    # variant needs to sanitize the inputs
+    arginds = list(toolz.partition(2, args))
+    arrays, inds = zip(*arginds)
+    if all(ind is None for ind in inds):
+        return {}, list(arrays), False
+    if all(ind == inds[0] for ind in inds) and all(
+        a.chunks == arrays[0].chunks for a in arrays
+    ):
+        return dict(zip(inds[0], arrays[0].chunks)), arrays, False
+
+    nameinds = []
+    blockdim_dict = dict()
+    for a, ind in arginds:
+        if ind is not None:
+            nameinds.append((a.name, ind))
+            blockdim_dict[a.name] = a.chunks
+        else:
+            nameinds.append((a, ind))
+
+    chunkss = broadcast_dimensions(nameinds, blockdim_dict, consolidate=common_blockdim)
+
+    arrays = []
+    changed = False
+    for a, i in arginds:
+        if i is None:
+            pass
+        else:
+            chunks = tuple(
+                (
+                    chunkss[j]
+                    if a.shape[n] > 1
+                    else (a.shape[n],) if not np.isnan(sum(chunkss[j])) else None
+                )
+                for n, j in enumerate(i)
+            )
+            if chunks != a.chunks and all(a.chunks):
+                a = a.rechunk(chunks)
+                changed = True
+            else:
+                pass
+        arrays.extend([a, i])
+    return chunkss, arrays, changed
diff --git a/dask/array/_array_expr/_rechunk.py b/dask/array/_array_expr/_rechunk.py
new file mode 100644
index 000000000..f46f09c2e
--- /dev/null
+++ b/dask/array/_array_expr/_rechunk.py
@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+import itertools
+import operator
+
+import numpy as np
+import toolz
+
+from dask.array._array_expr._expr import ArrayExpr
+from dask.array.core import concatenate3
+from dask.array.rechunk import (
+    _balance_chunksizes,
+    _choose_rechunk_method,
+    _validate_rechunk,
+    intersect_chunks,
+    normalize_chunks,
+    plan_rechunk,
+    tokenize,
+    validate_axis,
+)
+from dask.utils import cached_property
+
+
+class Rechunk(ArrayExpr):
+    _parameters = [
+        "array",
+        "_chunks",
+        "threshold",
+        "block_size_limit",
+        "balance",
+        "method",
+    ]
+
+    _defaults = {
+        "_chunks": "auto",
+        "threshold": None,
+        "block_size_limit": None,
+        "balance": None,
+        "method": None,
+    }
+
+    @property
+    def _meta(self):
+        return self.array._meta
+
+    @property
+    def _name(self):
+        return "rechunk-merge-" + tokenize(*self.operands)
+
+    @cached_property
+    def chunks(self):
+        x = self.array
+        chunks = self.operand("_chunks")
+
+        # don't rechunk if array is empty
+        if x.ndim > 0 and all(s == 0 for s in x.shape):
+            return x.chunks
+
+        if isinstance(chunks, dict):
+            chunks = {validate_axis(c, x.ndim): v for c, v in chunks.items()}
+            for i in range(x.ndim):
+                if i not in chunks:
+                    chunks[i] = x.chunks[i]
+                elif chunks[i] is None:
+                    chunks[i] = x.chunks[i]
+        if isinstance(chunks, (tuple, list)):
+            chunks = tuple(
+                lc if lc is not None else rc for lc, rc in zip(chunks, x.chunks)
+            )
+        chunks = normalize_chunks(
+            chunks,
+            x.shape,
+            limit=self.block_size_limit,
+            dtype=x.dtype,
+            previous_chunks=x.chunks,
+        )
+
+        if not len(chunks) == x.ndim:
+            raise ValueError("Provided chunks are not consistent with shape")
+
+        if self.balance:
+            chunks = tuple(_balance_chunksizes(chunk) for chunk in chunks)
+
+        _validate_rechunk(x.chunks, chunks)
+
+        return chunks
+
+    def _lower(self):
+
+        if not self.balance and (self.chunks == self.array.chunks):
+            return self.array
+
+        method = self.method or _choose_rechunk_method(
+            self.array.chunks, self.chunks, threshold=self.threshold
+        )
+        if method == "p2p":
+            raise NotImplementedError
+        else:
+            return TasksRechunk(
+                self.array, self.chunks, self.threshold, self.block_size_limit
+            )
+
+
+class TasksRechunk(Rechunk):
+    _parameters = ["array", "_chunks", "threshold", "block_size_limit"]
+
+    @cached_property
+    def chunks(self):
+        return self.operand("_chunks")
+
+    def _lower(self):
+        return
+
+    def _layer(self):
+        steps = plan_rechunk(
+            self.array.chunks,
+            self.chunks,
+            self.array.dtype.itemsize,
+            self.threshold,
+            self.block_size_limit,
+        )
+        name = self.array.name
+        old_chunks = self.array.chunks
+        layers = []
+        for i, c in enumerate(steps):
+            level = len(steps) - i - 1
+            name, old_chunks, layer = _compute_rechunk(
+                name, old_chunks, c, level, self.name
+            )
+            layers.append(layer)
+
+        return toolz.merge(*layers)
+
+
+def _compute_rechunk(old_name, old_chunks, chunks, level, name):
+    """Compute the rechunk of *x* to the given *chunks*."""
+    # TODO: redo this logic
+    # if x.size == 0:
+    #     # Special case for empty array, as the algorithm below does not behave correctly
+    #     return empty(x.shape, chunks=chunks, dtype=x.dtype)
+
+    ndim = len(old_chunks)
+    crossed = intersect_chunks(old_chunks, chunks)
+    x2 = dict()
+    intermediates = dict()
+    # token = tokenize(old_name, chunks)
+    if level != 0:
+        merge_name = name.replace("rechunk-merge-", f"rechunk-merge-{level}-")
+        split_name = name.replace("rechunk-merge-", f"rechunk-split-{level}-")
+    else:
+        merge_name = name.replace("rechunk-merge-", "rechunk-merge-")
+        split_name = name.replace("rechunk-merge-", "rechunk-split-")
+    split_name_suffixes = itertools.count()
+
+    # Pre-allocate old block references, to allow re-use and reduce the
+    # graph's memory footprint a bit.
+    old_blocks = np.empty([len(c) for c in old_chunks], dtype="O")
+    for index in np.ndindex(old_blocks.shape):
+        old_blocks[index] = (old_name,) + index
+
+    # Iterate over all new blocks
+    new_index = itertools.product(*(range(len(c)) for c in chunks))
+
+    for new_idx, cross1 in zip(new_index, crossed):
+        key = (merge_name,) + new_idx
+        old_block_indices = [[cr[i][0] for cr in cross1] for i in range(ndim)]
+        subdims1 = [len(set(old_block_indices[i])) for i in range(ndim)]
+
+        rec_cat_arg = np.empty(subdims1, dtype="O")
+        rec_cat_arg_flat = rec_cat_arg.flat
+
+        # Iterate over the old blocks required to build the new block
+        for rec_cat_index, ind_slices in enumerate(cross1):
+            old_block_index, slices = zip(*ind_slices)
+            name = (split_name, next(split_name_suffixes))
+            old_index = old_blocks[old_block_index][1:]
+            if all(
+                slc.start == 0 and slc.stop == old_chunks[i][ind]
+                for i, (slc, ind) in enumerate(zip(slices, old_index))
+            ):
+                rec_cat_arg_flat[rec_cat_index] = old_blocks[old_block_index]
+            else:
+                intermediates[name] = (
+                    operator.getitem,
+                    old_blocks[old_block_index],
+                    slices,
+                )
+                rec_cat_arg_flat[rec_cat_index] = name
+
+        assert rec_cat_index == rec_cat_arg.size - 1
+
+        # New block is formed by concatenation of sliced old blocks
+        if all(d == 1 for d in rec_cat_arg.shape):
+            x2[key] = rec_cat_arg.flat[0]
+        else:
+            x2[key] = (concatenate3, rec_cat_arg.tolist())
+
+    del old_blocks, new_index
+
+    return name, chunks, {**x2, **intermediates}
diff --git a/dask/array/_array_expr/_reductions.py b/dask/array/_array_expr/_reductions.py
new file mode 100644
index 000000000..2b65f5050
--- /dev/null
+++ b/dask/array/_array_expr/_reductions.py
@@ -0,0 +1,360 @@
+from __future__ import annotations
+
+import builtins
+import math
+from functools import partial
+from itertools import product
+from numbers import Integral
+
+from tlz import compose, get, partition_all
+
+from dask import config
+from dask._collections import new_collection
+from dask.array._array_expr._expr import ArrayExpr
+from dask.array._array_expr._utils import compute_meta
+from dask.array.core import _concatenate2
+from dask.array.utils import is_arraylike, validate_axis
+from dask.blockwise import lol_tuples
+from dask.tokenize import _tokenize_deterministic
+from dask.utils import cached_property, funcname, getargspec, is_series_like
+
+
+# TODO(expr): This needs somethign like what we have in DataFrame land with ACA
+# Way too many expressions that we are calling directly that should come from a
+# lower step.
+def reduction(
+    x,
+    chunk,
+    aggregate,
+    axis=None,
+    keepdims=False,
+    dtype=None,
+    split_every=None,
+    combine=None,
+    name=None,
+    out=None,
+    concatenate=True,
+    output_size=1,
+    meta=None,
+    weights=None,
+):
+    """General version of reductions
+
+    Parameters
+    ----------
+    x: Array
+        Data being reduced along one or more axes
+    chunk: callable(x_chunk, [weights_chunk=None], axis, keepdims)
+        First function to be executed when resolving the dask graph.
+        This function is applied in parallel to all original chunks of x.
+        See below for function parameters.
+    combine: callable(x_chunk, axis, keepdims), optional
+        Function used for intermediate recursive aggregation (see
+        split_every below). If omitted, it defaults to aggregate.
+        If the reduction can be performed in less than 3 steps, it will not
+        be invoked at all.
+    aggregate: callable(x_chunk, axis, keepdims)
+        Last function to be executed when resolving the dask graph,
+        producing the final output. It is always invoked, even when the reduced
+        Array counts a single chunk along the reduced axes.
+    axis: int or sequence of ints, optional
+        Axis or axes to aggregate upon. If omitted, aggregate along all axes.
+    keepdims: boolean, optional
+        Whether the reduction function should preserve the reduced axes,
+        leaving them at size ``output_size``, or remove them.
+    dtype: np.dtype
+        data type of output. This argument was previously optional, but
+        leaving as ``None`` will now raise an exception.
+    split_every: int >= 2 or dict(axis: int), optional
+        Determines the depth of the recursive aggregation. If set to or more
+        than the number of input chunks, the aggregation will be performed in
+        two steps, one ``chunk`` function per input chunk and a single
+        ``aggregate`` function at the end. If set to less than that, an
+        intermediate ``combine`` function will be used, so that any one
+        ``combine`` or ``aggregate`` function has no more than ``split_every``
+        inputs. The depth of the aggregation graph will be
+        :math:`log_{split_every}(input chunks along reduced axes)`. Setting to
+        a low value can reduce cache size and network transfers, at the cost of
+        more CPU and a larger dask graph.
+
+        Omit to let dask heuristically decide a good default. A default can
+        also be set globally with the ``split_every`` key in
+        :mod:`dask.config`.
+    name: str, optional
+        Prefix of the keys of the intermediate and output nodes. If omitted it
+        defaults to the function names.
+    out: Array, optional
+        Another dask array whose contents will be replaced. Omit to create a
+        new one. Note that, unlike in numpy, this setting gives no performance
+        benefits whatsoever, but can still be useful  if one needs to preserve
+        the references to a previously existing Array.
+    concatenate: bool, optional
+        If True (the default), the outputs of the ``chunk``/``combine``
+        functions are concatenated into a single np.array before being passed
+        to the ``combine``/``aggregate`` functions. If False, the input of
+        ``combine`` and ``aggregate`` will be either a list of the raw outputs
+        of the previous step or a single output, and the function will have to
+        concatenate it itself. It can be useful to set this to False if the
+        chunk and/or combine steps do not produce np.arrays.
+    output_size: int >= 1, optional
+        Size of the output of the ``aggregate`` function along the reduced
+        axes. Ignored if keepdims is False.
+    weights : array_like, optional
+        Weights to be used in the reduction of `x`. Will be
+        automatically broadcast to the shape of `x`, and so must have
+        a compatible shape. For instance, if `x` has shape ``(3, 4)``
+        then acceptable shapes for `weights` are ``(3, 4)``, ``(4,)``,
+        ``(3, 1)``, ``(1, 1)``, ``(1)``, and ``()``.
+
+    Returns
+    -------
+    dask array
+
+    **Function Parameters**
+
+    x_chunk: numpy.ndarray
+        Individual input chunk. For ``chunk`` functions, it is one of the
+        original chunks of x. For ``combine`` and ``aggregate`` functions, it's
+        the concatenation of the outputs produced by the previous ``chunk`` or
+        ``combine`` functions. If concatenate=False, it's a list of the raw
+        outputs from the previous functions.
+    weights_chunk: numpy.ndarray, optional
+        Only applicable to the ``chunk`` function. Weights, with the
+        same shape as `x_chunk`, to be applied during the reduction of
+        the individual input chunk. If ``weights`` have not been
+        provided then the function may omit this parameter. When
+        `weights_chunk` is included then it must occur immediately
+        after the `x_chunk` parameter, and must also have a default
+        value for cases when ``weights`` are not provided.
+    axis: tuple
+        Normalized list of axes to reduce upon, e.g. ``(0, )``
+        Scalar, negative, and None axes have been normalized away.
+        Note that some numpy reduction functions cannot reduce along multiple
+        axes at once and strictly require an int in input. Such functions have
+        to be wrapped to cope.
+    keepdims: bool
+        Whether the reduction function should preserve the reduced axes or
+        remove them.
+
+    """
+    if output_size != 1:
+        # TODO(expr-soon): i am not convinced that this is actually used. We don't have a single
+        # test for this as far as I can tell. If we want to keep this, we will most likely need
+        # different arguments for chunk and aggregate to handle differing behaviors
+        raise NotImplementedError("output_size != 1 is not yet supported")
+    if axis is None:
+        axis = tuple(range(x.ndim))
+    if isinstance(axis, Integral):
+        axis = (axis,)
+    axis = validate_axis(axis, x.ndim)
+
+    if dtype is None:
+        raise ValueError("Must specify dtype")
+    if "dtype" in getargspec(chunk).args:
+        chunk = partial(chunk, dtype=dtype)
+    if "dtype" in getargspec(aggregate).args:
+        aggregate = partial(aggregate, dtype=dtype)
+    if is_series_like(x):
+        x = x.values
+
+    # Map chunk across all blocks
+    inds = tuple(range(x.ndim))
+
+    args = (x.expr, inds)
+
+    if weights is not None:
+        # TODO(expr-soon): Needs more IO Stuff
+        raise NotImplementedError("Weights are not yet supported")
+
+    # The dtype of `tmp` doesn't actually matter, and may be incorrect.
+    tmp = blockwise(
+        chunk, inds, *args, axis=axis, keepdims=True, token=name, dtype=dtype or float
+    )
+    if meta is None and hasattr(x, "_meta"):
+        try:
+            reduced_meta = compute_meta(
+                chunk, x.dtype, x._meta, axis=axis, keepdims=True, computing_meta=True
+            )
+        except TypeError:
+            reduced_meta = compute_meta(
+                chunk, x.dtype, x._meta, axis=axis, keepdims=True
+            )
+        except ValueError:
+            reduced_meta = None
+    else:
+        reduced_meta = None
+
+    result = _tree_reduce(
+        tmp,
+        aggregate,
+        axis,
+        keepdims,
+        dtype,
+        split_every,
+        combine,
+        name=name,
+        concatenate=concatenate,
+        reduced_meta=reduced_meta if reduced_meta is not None else meta,
+    )
+    return result
+
+
+def _tree_reduce(
+    x,
+    aggregate,
+    axis,
+    keepdims,
+    dtype,
+    split_every=None,
+    combine=None,
+    name=None,
+    concatenate=True,
+    reduced_meta=None,
+):
+    """Perform the tree reduction step of a reduction.
+
+    Lower level, users should use ``reduction`` or ``arg_reduction`` directly.
+    """
+    # Normalize split_every
+    split_every = split_every or config.get("split_every", 16)
+    if isinstance(split_every, dict):
+        split_every = {k: split_every.get(k, 2) for k in axis}
+    elif isinstance(split_every, Integral):
+        n = builtins.max(int(split_every ** (1 / (len(axis) or 1))), 2)
+        split_every = dict.fromkeys(axis, n)
+    else:
+        raise ValueError("split_every must be a int or a dict")
+
+    # Reduce across intermediates
+    depth = 1
+    for i, n in enumerate(x.numblocks):
+        if i in split_every and split_every[i] != 1:
+            depth = int(builtins.max(depth, math.ceil(math.log(n, split_every[i]))))
+    func = partial(combine or aggregate, axis=axis, keepdims=True)
+    if concatenate:
+        func = compose(func, partial(_concatenate2, axes=sorted(axis)))
+    for _ in range(depth - 1):
+        x = PartialReduce(
+            x,
+            func,
+            split_every,
+            True,
+            dtype=dtype,
+            name=(name or funcname(combine or aggregate)) + "-partial",
+            reduced_meta=reduced_meta,
+        )
+    func = partial(aggregate, axis=axis, keepdims=keepdims)
+    if concatenate:
+        func = compose(func, partial(_concatenate2, axes=sorted(axis)))
+    return new_collection(
+        PartialReduce(
+            x,
+            func,
+            split_every,
+            keepdims=keepdims,
+            dtype=dtype,
+            name=(name or funcname(aggregate)) + "-aggregate",
+            reduced_meta=reduced_meta,
+        )
+    )
+
+
+class PartialReduce(ArrayExpr):
+    _parameters = [
+        "array",
+        "func",
+        "split_every",
+        "keepdims",
+        "dtype",
+        "name",
+        "reduced_meta",
+    ]
+    _defaults = {
+        "keepdims": False,
+        "dtype": None,
+        "name": None,
+        "reduced_meta": None,
+    }
+
+    @cached_property
+    def _name(self):
+        return (
+            (self.operand("name") or funcname(self.func))
+            + "-"
+            + _tokenize_deterministic(
+                self.func, self.array, self.split_every, self.keepdims, self.dtype
+            )
+        )
+
+    @cached_property
+    def chunks(self):
+        chunks = [
+            (
+                tuple(1 for p in partition_all(self.split_every[i], c))
+                if i in self.split_every
+                else c
+            )
+            for (i, c) in enumerate(self.array.chunks)
+        ]
+
+        if not self.keepdims:
+            out_axis = [i for i in range(self.array.ndim) if i not in self.split_every]
+            getter = lambda k: get(out_axis, k)
+            chunks = list(getter(chunks))
+        return tuple(chunks)
+
+    def _layer(self):
+        x = self.array
+        parts = [
+            list(partition_all(self.split_every.get(i, 1), range(n)))
+            for (i, n) in enumerate(x.numblocks)
+        ]
+        keys = product(*map(range, map(len, parts)))
+        if not self.keepdims:
+            out_axis = [i for i in range(x.ndim) if i not in self.split_every]
+            getter = lambda k: get(out_axis, k)
+            keys = map(getter, keys)
+        dsk = {}
+        for k, p in zip(keys, product(*parts)):
+            free = {
+                i: j[0]
+                for (i, j) in enumerate(p)
+                if len(j) == 1 and i not in self.split_every
+            }
+            dummy = dict(i for i in enumerate(p) if i[0] in self.split_every)
+            g = lol_tuples((x.name,), range(x.ndim), free, dummy)
+            dsk[(self._name,) + k] = (self.func, g)
+
+        return dsk
+
+    @property
+    def _meta(self):
+        meta = self.array._meta
+        if self.reduced_meta is not None:
+            try:
+                meta = self.func(self.reduced_meta, computing_meta=True)
+            # no meta keyword argument exists for func, and it isn't required
+            except TypeError:
+                try:
+                    meta = self.func(self.reduced_meta)
+                except ValueError as e:
+                    # min/max functions have no identity, don't apply function to meta
+                    if "zero-size array to reduction operation" in str(e):
+                        meta = self.reduced_meta
+            # when no work can be computed on the empty array (e.g., func is a ufunc)
+            except ValueError:
+                pass
+
+        # some functions can't compute empty arrays (those for which reduced_meta
+        # fall into the ValueError exception) and we have to rely on reshaping
+        # the array according to len(out_chunks)
+        if is_arraylike(meta) and meta.ndim != len(self.chunks):
+            if len(self.chunks) == 0:
+                meta = meta.sum()
+            else:
+                meta = meta.reshape((0,) * len(self.chunks))
+
+        return meta
+
+
+from dask.array._array_expr._collection import blockwise
diff --git a/dask/array/_reductions_generic.py b/dask/array/_reductions_generic.py
new file mode 100644
index 000000000..3dee26537
--- /dev/null
+++ b/dask/array/_reductions_generic.py
@@ -0,0 +1,351 @@
+from __future__ import annotations
+
+import builtins
+import contextlib
+import math
+import warnings
+from functools import partial
+from itertools import product
+from numbers import Integral
+
+import numpy as np
+from tlz import compose, get, partition_all
+
+from dask import config
+from dask.array.blockwise import blockwise
+from dask.array.core import Array, _concatenate2, asanyarray, broadcast_to, handle_out
+from dask.array.numpy_compat import ComplexWarning
+from dask.array.utils import compute_meta, is_arraylike, validate_axis
+from dask.base import tokenize
+from dask.blockwise import lol_tuples
+from dask.highlevelgraph import HighLevelGraph
+from dask.utils import funcname, getargspec, is_series_like
+
+
+def reduction(
+    x,
+    chunk,
+    aggregate,
+    axis=None,
+    keepdims=False,
+    dtype=None,
+    split_every=None,
+    combine=None,
+    name=None,
+    out=None,
+    concatenate=True,
+    output_size=1,
+    meta=None,
+    weights=None,
+):
+    """General version of reductions
+
+    Parameters
+    ----------
+    x: Array
+        Data being reduced along one or more axes
+    chunk: callable(x_chunk, [weights_chunk=None], axis, keepdims)
+        First function to be executed when resolving the dask graph.
+        This function is applied in parallel to all original chunks of x.
+        See below for function parameters.
+    combine: callable(x_chunk, axis, keepdims), optional
+        Function used for intermediate recursive aggregation (see
+        split_every below). If omitted, it defaults to aggregate.
+        If the reduction can be performed in less than 3 steps, it will not
+        be invoked at all.
+    aggregate: callable(x_chunk, axis, keepdims)
+        Last function to be executed when resolving the dask graph,
+        producing the final output. It is always invoked, even when the reduced
+        Array counts a single chunk along the reduced axes.
+    axis: int or sequence of ints, optional
+        Axis or axes to aggregate upon. If omitted, aggregate along all axes.
+    keepdims: boolean, optional
+        Whether the reduction function should preserve the reduced axes,
+        leaving them at size ``output_size``, or remove them.
+    dtype: np.dtype
+        data type of output. This argument was previously optional, but
+        leaving as ``None`` will now raise an exception.
+    split_every: int >= 2 or dict(axis: int), optional
+        Determines the depth of the recursive aggregation. If set to or more
+        than the number of input chunks, the aggregation will be performed in
+        two steps, one ``chunk`` function per input chunk and a single
+        ``aggregate`` function at the end. If set to less than that, an
+        intermediate ``combine`` function will be used, so that any one
+        ``combine`` or ``aggregate`` function has no more than ``split_every``
+        inputs. The depth of the aggregation graph will be
+        :math:`log_{split_every}(input chunks along reduced axes)`. Setting to
+        a low value can reduce cache size and network transfers, at the cost of
+        more CPU and a larger dask graph.
+
+        Omit to let dask heuristically decide a good default. A default can
+        also be set globally with the ``split_every`` key in
+        :mod:`dask.config`.
+    name: str, optional
+        Prefix of the keys of the intermediate and output nodes. If omitted it
+        defaults to the function names.
+    out: Array, optional
+        Another dask array whose contents will be replaced. Omit to create a
+        new one. Note that, unlike in numpy, this setting gives no performance
+        benefits whatsoever, but can still be useful  if one needs to preserve
+        the references to a previously existing Array.
+    concatenate: bool, optional
+        If True (the default), the outputs of the ``chunk``/``combine``
+        functions are concatenated into a single np.array before being passed
+        to the ``combine``/``aggregate`` functions. If False, the input of
+        ``combine`` and ``aggregate`` will be either a list of the raw outputs
+        of the previous step or a single output, and the function will have to
+        concatenate it itself. It can be useful to set this to False if the
+        chunk and/or combine steps do not produce np.arrays.
+    output_size: int >= 1, optional
+        Size of the output of the ``aggregate`` function along the reduced
+        axes. Ignored if keepdims is False.
+    weights : array_like, optional
+        Weights to be used in the reduction of `x`. Will be
+        automatically broadcast to the shape of `x`, and so must have
+        a compatible shape. For instance, if `x` has shape ``(3, 4)``
+        then acceptable shapes for `weights` are ``(3, 4)``, ``(4,)``,
+        ``(3, 1)``, ``(1, 1)``, ``(1)``, and ``()``.
+
+    Returns
+    -------
+    dask array
+
+    **Function Parameters**
+
+    x_chunk: numpy.ndarray
+        Individual input chunk. For ``chunk`` functions, it is one of the
+        original chunks of x. For ``combine`` and ``aggregate`` functions, it's
+        the concatenation of the outputs produced by the previous ``chunk`` or
+        ``combine`` functions. If concatenate=False, it's a list of the raw
+        outputs from the previous functions.
+    weights_chunk: numpy.ndarray, optional
+        Only applicable to the ``chunk`` function. Weights, with the
+        same shape as `x_chunk`, to be applied during the reduction of
+        the individual input chunk. If ``weights`` have not been
+        provided then the function may omit this parameter. When
+        `weights_chunk` is included then it must occur immediately
+        after the `x_chunk` parameter, and must also have a default
+        value for cases when ``weights`` are not provided.
+    axis: tuple
+        Normalized list of axes to reduce upon, e.g. ``(0, )``
+        Scalar, negative, and None axes have been normalized away.
+        Note that some numpy reduction functions cannot reduce along multiple
+        axes at once and strictly require an int in input. Such functions have
+        to be wrapped to cope.
+    keepdims: bool
+        Whether the reduction function should preserve the reduced axes or
+        remove them.
+
+    """
+    if axis is None:
+        axis = tuple(range(x.ndim))
+    if isinstance(axis, Integral):
+        axis = (axis,)
+    axis = validate_axis(axis, x.ndim)
+
+    if dtype is None:
+        raise ValueError("Must specify dtype")
+    if "dtype" in getargspec(chunk).args:
+        chunk = partial(chunk, dtype=dtype)
+    if "dtype" in getargspec(aggregate).args:
+        aggregate = partial(aggregate, dtype=dtype)
+    if is_series_like(x):
+        x = x.values
+
+    # Map chunk across all blocks
+    inds = tuple(range(x.ndim))
+
+    args = (x, inds)
+    if weights is not None:
+        # Broadcast weights to x and add to args
+        wgt = asanyarray(weights)
+        try:
+            wgt = broadcast_to(wgt, x.shape)
+        except ValueError:
+            raise ValueError(
+                f"Weights with shape {wgt.shape} are not broadcastable "
+                f"to x with shape {x.shape}"
+            )
+
+        args += (wgt, inds)
+
+    # The dtype of `tmp` doesn't actually matter, and may be incorrect.
+    tmp = blockwise(
+        chunk, inds, *args, axis=axis, keepdims=True, token=name, dtype=dtype or float
+    )
+    tmp._chunks = tuple(
+        (output_size,) * len(c) if i in axis else c for i, c in enumerate(tmp.chunks)
+    )
+
+    if meta is None and hasattr(x, "_meta"):
+        try:
+            reduced_meta = compute_meta(
+                chunk, x.dtype, x._meta, axis=axis, keepdims=True, computing_meta=True
+            )
+        except TypeError:
+            reduced_meta = compute_meta(
+                chunk, x.dtype, x._meta, axis=axis, keepdims=True
+            )
+        except ValueError:
+            pass
+    else:
+        reduced_meta = None
+
+    result = _tree_reduce(
+        tmp,
+        aggregate,
+        axis,
+        keepdims,
+        dtype,
+        split_every,
+        combine,
+        name=name,
+        concatenate=concatenate,
+        reduced_meta=reduced_meta,
+    )
+    if keepdims and output_size != 1:
+        result._chunks = tuple(
+            (output_size,) if i in axis else c for i, c in enumerate(tmp.chunks)
+        )
+    if meta is not None:
+        result._meta = meta
+    return handle_out(out, result)
+
+
+def _tree_reduce(
+    x,
+    aggregate,
+    axis,
+    keepdims,
+    dtype,
+    split_every=None,
+    combine=None,
+    name=None,
+    concatenate=True,
+    reduced_meta=None,
+):
+    """Perform the tree reduction step of a reduction.
+
+    Lower level, users should use ``reduction`` or ``arg_reduction`` directly.
+    """
+    # Normalize split_every
+    split_every = split_every or config.get("split_every", 4)
+    if isinstance(split_every, dict):
+        split_every = {k: split_every.get(k, 2) for k in axis}
+    elif isinstance(split_every, Integral):
+        n = builtins.max(int(split_every ** (1 / (len(axis) or 1))), 2)
+        split_every = dict.fromkeys(axis, n)
+    else:
+        raise ValueError("split_every must be a int or a dict")
+
+    # Reduce across intermediates
+    depth = 1
+    for i, n in enumerate(x.numblocks):
+        if i in split_every and split_every[i] != 1:
+            depth = int(builtins.max(depth, math.ceil(math.log(n, split_every[i]))))
+    func = partial(combine or aggregate, axis=axis, keepdims=True)
+    if concatenate:
+        func = compose(func, partial(_concatenate2, axes=sorted(axis)))
+    for _ in range(depth - 1):
+        x = partial_reduce(
+            func,
+            x,
+            split_every,
+            True,
+            dtype=dtype,
+            name=(name or funcname(combine or aggregate)) + "-partial",
+            reduced_meta=reduced_meta,
+        )
+    func = partial(aggregate, axis=axis, keepdims=keepdims)
+    if concatenate:
+        func = compose(func, partial(_concatenate2, axes=sorted(axis)))
+    return partial_reduce(
+        func,
+        x,
+        split_every,
+        keepdims=keepdims,
+        dtype=dtype,
+        name=(name or funcname(aggregate)) + "-aggregate",
+        reduced_meta=reduced_meta,
+    )
+
+
+def partial_reduce(
+    func, x, split_every, keepdims=False, dtype=None, name=None, reduced_meta=None
+):
+    """Partial reduction across multiple axes.
+
+    Parameters
+    ----------
+    func : function
+    x : Array
+    split_every : dict
+        Maximum reduction block sizes in each dimension.
+
+    Examples
+    --------
+    Reduce across axis 0 and 2, merging a maximum of 1 block in the 0th
+    dimension, and 3 blocks in the 2nd dimension:
+
+    >>> partial_reduce(np.min, x, {0: 1, 2: 3})    # doctest: +SKIP
+    """
+    name = (
+        (name or funcname(func)) + "-" + tokenize(func, x, split_every, keepdims, dtype)
+    )
+    parts = [
+        list(partition_all(split_every.get(i, 1), range(n)))
+        for (i, n) in enumerate(x.numblocks)
+    ]
+    keys = product(*map(range, map(len, parts)))
+    out_chunks = [
+        tuple(1 for p in partition_all(split_every[i], c)) if i in split_every else c
+        for (i, c) in enumerate(x.chunks)
+    ]
+    if not keepdims:
+        out_axis = [i for i in range(x.ndim) if i not in split_every]
+        getter = lambda k: get(out_axis, k)
+        keys = map(getter, keys)
+        out_chunks = list(getter(out_chunks))
+    dsk = {}
+    for k, p in zip(keys, product(*parts)):
+        free = {
+            i: j[0] for (i, j) in enumerate(p) if len(j) == 1 and i not in split_every
+        }
+        dummy = dict(i for i in enumerate(p) if i[0] in split_every)
+        g = lol_tuples((x.name,), range(x.ndim), free, dummy)
+        dsk[(name,) + k] = (func, g)
+    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[x])
+
+    meta = x._meta
+    if reduced_meta is not None:
+        try:
+            meta = func(reduced_meta, computing_meta=True)
+        # no meta keyword argument exists for func, and it isn't required
+        except TypeError:
+            try:
+                meta = func(reduced_meta)
+            except ValueError as e:
+                # min/max functions have no identity, don't apply function to meta
+                if "zero-size array to reduction operation" in str(e):
+                    meta = reduced_meta
+        # when no work can be computed on the empty array (e.g., func is a ufunc)
+        except ValueError:
+            pass
+
+    # some functions can't compute empty arrays (those for which reduced_meta
+    # fall into the ValueError exception) and we have to rely on reshaping
+    # the array according to len(out_chunks)
+    if is_arraylike(meta) and meta.ndim != len(out_chunks):
+        if len(out_chunks) == 0:
+            meta = meta.sum()
+        else:
+            meta = meta.reshape((0,) * len(out_chunks))
+
+    if np.isscalar(meta):
+        return Array(graph, name, out_chunks, dtype=dtype)
+    else:
+        with contextlib.suppress(AttributeError), warnings.catch_warnings():
+            if name.startswith("var") or name.startswith("moment"):
+                warnings.simplefilter("ignore", ComplexWarning)
+            meta = meta.astype(dtype)
+        return Array(graph, name, out_chunks, meta=meta)
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
index 8f668815f..27aafc3f8 100644
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import builtins
-import contextlib
 import math
 import operator
 import warnings
@@ -12,43 +11,36 @@ from numbers import Integral, Number
 from operator import mul
 
 import numpy as np
-from tlz import accumulate, compose, drop, get, partition_all, pluck
+from tlz import accumulate, drop, pluck
 
-from dask import config
+import dask.array as da
 from dask.array import chunk
-from dask.array.blockwise import blockwise
 from dask.array.core import (
     Array,
     _concatenate2,
-    asanyarray,
-    broadcast_to,
     handle_out,
     implements,
     unknown_chunk_message,
 )
 from dask.array.creation import arange, diagonal
 from dask.array.dispatch import divide_lookup, nannumel_lookup, numel_lookup
-from dask.array.numpy_compat import NUMPY_GE_200, ComplexWarning
+from dask.array.numpy_compat import NUMPY_GE_200
 from dask.array.utils import (
     array_safe,
     asarray_safe,
-    compute_meta,
     is_arraylike,
     meta_from_array,
     validate_axis,
 )
 from dask.array.wrap import ones, zeros
 from dask.base import tokenize
-from dask.blockwise import lol_tuples
 from dask.highlevelgraph import HighLevelGraph
-from dask.utils import (
-    apply,
-    deepmap,
-    derived_from,
-    funcname,
-    getargspec,
-    is_series_like,
-)
+from dask.utils import apply, deepmap, derived_from
+
+if da._array_expr_enabled():
+    from dask.array._array_expr import _tree_reduce, reduction
+else:
+    from dask.array._reductions_generic import _tree_reduce, reduction
 
 
 def divide(a, b, dtype=None):
@@ -65,335 +57,6 @@ def nannumel(x, **kwargs):
     return nannumel_lookup(x, **kwargs)
 
 
-def reduction(
-    x,
-    chunk,
-    aggregate,
-    axis=None,
-    keepdims=False,
-    dtype=None,
-    split_every=None,
-    combine=None,
-    name=None,
-    out=None,
-    concatenate=True,
-    output_size=1,
-    meta=None,
-    weights=None,
-):
-    """General version of reductions
-
-    Parameters
-    ----------
-    x: Array
-        Data being reduced along one or more axes
-    chunk: callable(x_chunk, [weights_chunk=None], axis, keepdims)
-        First function to be executed when resolving the dask graph.
-        This function is applied in parallel to all original chunks of x.
-        See below for function parameters.
-    combine: callable(x_chunk, axis, keepdims), optional
-        Function used for intermediate recursive aggregation (see
-        split_every below). If omitted, it defaults to aggregate.
-        If the reduction can be performed in less than 3 steps, it will not
-        be invoked at all.
-    aggregate: callable(x_chunk, axis, keepdims)
-        Last function to be executed when resolving the dask graph,
-        producing the final output. It is always invoked, even when the reduced
-        Array counts a single chunk along the reduced axes.
-    axis: int or sequence of ints, optional
-        Axis or axes to aggregate upon. If omitted, aggregate along all axes.
-    keepdims: boolean, optional
-        Whether the reduction function should preserve the reduced axes,
-        leaving them at size ``output_size``, or remove them.
-    dtype: np.dtype
-        data type of output. This argument was previously optional, but
-        leaving as ``None`` will now raise an exception.
-    split_every: int >= 2 or dict(axis: int), optional
-        Determines the depth of the recursive aggregation. If set to or more
-        than the number of input chunks, the aggregation will be performed in
-        two steps, one ``chunk`` function per input chunk and a single
-        ``aggregate`` function at the end. If set to less than that, an
-        intermediate ``combine`` function will be used, so that any one
-        ``combine`` or ``aggregate`` function has no more than ``split_every``
-        inputs. The depth of the aggregation graph will be
-        :math:`log_{split_every}(input chunks along reduced axes)`. Setting to
-        a low value can reduce cache size and network transfers, at the cost of
-        more CPU and a larger dask graph.
-
-        Omit to let dask heuristically decide a good default. A default can
-        also be set globally with the ``split_every`` key in
-        :mod:`dask.config`.
-    name: str, optional
-        Prefix of the keys of the intermediate and output nodes. If omitted it
-        defaults to the function names.
-    out: Array, optional
-        Another dask array whose contents will be replaced. Omit to create a
-        new one. Note that, unlike in numpy, this setting gives no performance
-        benefits whatsoever, but can still be useful  if one needs to preserve
-        the references to a previously existing Array.
-    concatenate: bool, optional
-        If True (the default), the outputs of the ``chunk``/``combine``
-        functions are concatenated into a single np.array before being passed
-        to the ``combine``/``aggregate`` functions. If False, the input of
-        ``combine`` and ``aggregate`` will be either a list of the raw outputs
-        of the previous step or a single output, and the function will have to
-        concatenate it itself. It can be useful to set this to False if the
-        chunk and/or combine steps do not produce np.arrays.
-    output_size: int >= 1, optional
-        Size of the output of the ``aggregate`` function along the reduced
-        axes. Ignored if keepdims is False.
-    weights : array_like, optional
-        Weights to be used in the reduction of `x`. Will be
-        automatically broadcast to the shape of `x`, and so must have
-        a compatible shape. For instance, if `x` has shape ``(3, 4)``
-        then acceptable shapes for `weights` are ``(3, 4)``, ``(4,)``,
-        ``(3, 1)``, ``(1, 1)``, ``(1)``, and ``()``.
-
-    Returns
-    -------
-    dask array
-
-    **Function Parameters**
-
-    x_chunk: numpy.ndarray
-        Individual input chunk. For ``chunk`` functions, it is one of the
-        original chunks of x. For ``combine`` and ``aggregate`` functions, it's
-        the concatenation of the outputs produced by the previous ``chunk`` or
-        ``combine`` functions. If concatenate=False, it's a list of the raw
-        outputs from the previous functions.
-    weights_chunk: numpy.ndarray, optional
-        Only applicable to the ``chunk`` function. Weights, with the
-        same shape as `x_chunk`, to be applied during the reduction of
-        the individual input chunk. If ``weights`` have not been
-        provided then the function may omit this parameter. When
-        `weights_chunk` is included then it must occur immediately
-        after the `x_chunk` parameter, and must also have a default
-        value for cases when ``weights`` are not provided.
-    axis: tuple
-        Normalized list of axes to reduce upon, e.g. ``(0, )``
-        Scalar, negative, and None axes have been normalized away.
-        Note that some numpy reduction functions cannot reduce along multiple
-        axes at once and strictly require an int in input. Such functions have
-        to be wrapped to cope.
-    keepdims: bool
-        Whether the reduction function should preserve the reduced axes or
-        remove them.
-
-    """
-    if axis is None:
-        axis = tuple(range(x.ndim))
-    if isinstance(axis, Integral):
-        axis = (axis,)
-    axis = validate_axis(axis, x.ndim)
-
-    if dtype is None:
-        raise ValueError("Must specify dtype")
-    if "dtype" in getargspec(chunk).args:
-        chunk = partial(chunk, dtype=dtype)
-    if "dtype" in getargspec(aggregate).args:
-        aggregate = partial(aggregate, dtype=dtype)
-    if is_series_like(x):
-        x = x.values
-
-    # Map chunk across all blocks
-    inds = tuple(range(x.ndim))
-
-    args = (x, inds)
-    if weights is not None:
-        # Broadcast weights to x and add to args
-        wgt = asanyarray(weights)
-        try:
-            wgt = broadcast_to(wgt, x.shape)
-        except ValueError:
-            raise ValueError(
-                f"Weights with shape {wgt.shape} are not broadcastable "
-                f"to x with shape {x.shape}"
-            )
-
-        args += (wgt, inds)
-
-    # The dtype of `tmp` doesn't actually matter, and may be incorrect.
-    tmp = blockwise(
-        chunk, inds, *args, axis=axis, keepdims=True, token=name, dtype=dtype or float
-    )
-    tmp._chunks = tuple(
-        (output_size,) * len(c) if i in axis else c for i, c in enumerate(tmp.chunks)
-    )
-
-    if meta is None and hasattr(x, "_meta"):
-        try:
-            reduced_meta = compute_meta(
-                chunk, x.dtype, x._meta, axis=axis, keepdims=True, computing_meta=True
-            )
-        except TypeError:
-            reduced_meta = compute_meta(
-                chunk, x.dtype, x._meta, axis=axis, keepdims=True
-            )
-        except ValueError:
-            pass
-    else:
-        reduced_meta = None
-
-    result = _tree_reduce(
-        tmp,
-        aggregate,
-        axis,
-        keepdims,
-        dtype,
-        split_every,
-        combine,
-        name=name,
-        concatenate=concatenate,
-        reduced_meta=reduced_meta,
-    )
-    if keepdims and output_size != 1:
-        result._chunks = tuple(
-            (output_size,) if i in axis else c for i, c in enumerate(tmp.chunks)
-        )
-    if meta is not None:
-        result._meta = meta
-    return handle_out(out, result)
-
-
-def _tree_reduce(
-    x,
-    aggregate,
-    axis,
-    keepdims,
-    dtype,
-    split_every=None,
-    combine=None,
-    name=None,
-    concatenate=True,
-    reduced_meta=None,
-):
-    """Perform the tree reduction step of a reduction.
-
-    Lower level, users should use ``reduction`` or ``arg_reduction`` directly.
-    """
-    # Normalize split_every
-    split_every = split_every or config.get("split_every", 4)
-    if isinstance(split_every, dict):
-        split_every = {k: split_every.get(k, 2) for k in axis}
-    elif isinstance(split_every, Integral):
-        n = builtins.max(int(split_every ** (1 / (len(axis) or 1))), 2)
-        split_every = dict.fromkeys(axis, n)
-    else:
-        raise ValueError("split_every must be a int or a dict")
-
-    # Reduce across intermediates
-    depth = 1
-    for i, n in enumerate(x.numblocks):
-        if i in split_every and split_every[i] != 1:
-            depth = int(builtins.max(depth, math.ceil(math.log(n, split_every[i]))))
-    func = partial(combine or aggregate, axis=axis, keepdims=True)
-    if concatenate:
-        func = compose(func, partial(_concatenate2, axes=sorted(axis)))
-    for _ in range(depth - 1):
-        x = partial_reduce(
-            func,
-            x,
-            split_every,
-            True,
-            dtype=dtype,
-            name=(name or funcname(combine or aggregate)) + "-partial",
-            reduced_meta=reduced_meta,
-        )
-    func = partial(aggregate, axis=axis, keepdims=keepdims)
-    if concatenate:
-        func = compose(func, partial(_concatenate2, axes=sorted(axis)))
-    return partial_reduce(
-        func,
-        x,
-        split_every,
-        keepdims=keepdims,
-        dtype=dtype,
-        name=(name or funcname(aggregate)) + "-aggregate",
-        reduced_meta=reduced_meta,
-    )
-
-
-def partial_reduce(
-    func, x, split_every, keepdims=False, dtype=None, name=None, reduced_meta=None
-):
-    """Partial reduction across multiple axes.
-
-    Parameters
-    ----------
-    func : function
-    x : Array
-    split_every : dict
-        Maximum reduction block sizes in each dimension.
-
-    Examples
-    --------
-    Reduce across axis 0 and 2, merging a maximum of 1 block in the 0th
-    dimension, and 3 blocks in the 2nd dimension:
-
-    >>> partial_reduce(np.min, x, {0: 1, 2: 3})    # doctest: +SKIP
-    """
-    name = (
-        (name or funcname(func)) + "-" + tokenize(func, x, split_every, keepdims, dtype)
-    )
-    parts = [
-        list(partition_all(split_every.get(i, 1), range(n)))
-        for (i, n) in enumerate(x.numblocks)
-    ]
-    keys = product(*map(range, map(len, parts)))
-    out_chunks = [
-        tuple(1 for p in partition_all(split_every[i], c)) if i in split_every else c
-        for (i, c) in enumerate(x.chunks)
-    ]
-    if not keepdims:
-        out_axis = [i for i in range(x.ndim) if i not in split_every]
-        getter = lambda k: get(out_axis, k)
-        keys = map(getter, keys)
-        out_chunks = list(getter(out_chunks))
-    dsk = {}
-    for k, p in zip(keys, product(*parts)):
-        free = {
-            i: j[0] for (i, j) in enumerate(p) if len(j) == 1 and i not in split_every
-        }
-        dummy = dict(i for i in enumerate(p) if i[0] in split_every)
-        g = lol_tuples((x.name,), range(x.ndim), free, dummy)
-        dsk[(name,) + k] = (func, g)
-    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[x])
-
-    meta = x._meta
-    if reduced_meta is not None:
-        try:
-            meta = func(reduced_meta, computing_meta=True)
-        # no meta keyword argument exists for func, and it isn't required
-        except TypeError:
-            try:
-                meta = func(reduced_meta)
-            except ValueError as e:
-                # min/max functions have no identity, don't apply function to meta
-                if "zero-size array to reduction operation" in str(e):
-                    meta = reduced_meta
-        # when no work can be computed on the empty array (e.g., func is a ufunc)
-        except ValueError:
-            pass
-
-    # some functions can't compute empty arrays (those for which reduced_meta
-    # fall into the ValueError exception) and we have to rely on reshaping
-    # the array according to len(out_chunks)
-    if is_arraylike(meta) and meta.ndim != len(out_chunks):
-        if len(out_chunks) == 0:
-            meta = meta.sum()
-        else:
-            meta = meta.reshape((0,) * len(out_chunks))
-
-    if np.isscalar(meta):
-        return Array(graph, name, out_chunks, dtype=dtype)
-    else:
-        with contextlib.suppress(AttributeError), warnings.catch_warnings():
-            if name.startswith("var") or name.startswith("moment"):
-                warnings.simplefilter("ignore", ComplexWarning)
-            meta = meta.astype(dtype)
-        return Array(graph, name, out_chunks, meta=meta)
-
-
 @derived_from(np)
 def sum(a, axis=None, dtype=None, keepdims=False, split_every=None, out=None):
     if dtype is None:


Test patch:
diff --git a/dask/array/_array_expr/tests/test_collection.py b/dask/array/_array_expr/tests/test_collection.py
index e2f5362d3..fa3581664 100644
--- a/dask/array/_array_expr/tests/test_collection.py
+++ b/dask/array/_array_expr/tests/test_collection.py
@@ -1,9 +1,13 @@
 from __future__ import annotations
 
+import operator
+
+import numpy as np
 import pytest
 
-import dask.array._array_expr as da
+import dask.array as da
 from dask.array import assert_eq
+from dask.array._array_expr._rechunk import Rechunk
 
 
 @pytest.fixture()
@@ -33,3 +37,62 @@ def test_arithmetic_ops(arr, op):
     result = getattr(arr, op)(2)
     expected = getattr(arr.compute(), op)(2)
     assert_eq(result, expected)
+
+
+@pytest.mark.array_expr
+def test_rechunk(arr):
+    result = arr.rechunk((7, 3))
+    expected = arr.compute()
+    assert_eq(result, expected)
+
+
+@pytest.mark.array_expr
+def test_blockwise():
+    x = da.random.random((10, 10), chunks=(5, 5))
+    z = da.blockwise(operator.add, "ij", x, "ij", 100, None, dtype=x.dtype)
+    assert_eq(z, x.compute() + 100)
+
+    x = da.random.random((10, 10), chunks=(5, 5))
+    z = da.blockwise(operator.add, "ij", x, "ij", x, "ij", dtype=x.dtype)
+    expr = z.expr.optimize()
+    assert len(list(expr.find_operations(Rechunk))) == 0
+    assert_eq(z, x.compute() * 2)
+
+    # align
+    x = da.random.random((10, 10), chunks=(5, 5))
+    y = da.random.random((10, 10), chunks=(7, 3))
+    z = da.blockwise(operator.add, "ij", x, "ij", y, "ij", dtype=x.dtype)
+    expr = z.expr.optimize()
+    assert len(list(expr.find_operations(Rechunk))) > 0
+    assert_eq(z, x.compute() + y.compute())
+
+
+@pytest.mark.parametrize("func", ["min", "max", "sum", "prod", "mean", "any", "all"])
+def test_reductions(arr, func):
+    # var and std need __array_function__
+    result = getattr(arr, func)(axis=0)
+    expected = getattr(arr.compute(), func)(axis=0)
+    assert_eq(result, expected)
+
+
+@pytest.mark.parametrize(
+    "func",
+    [
+        "sum",
+        "mean",
+        "any",
+        "all",
+        "max",
+        "min",
+        "nanmin",
+        "nanmax",
+        "nanmean",
+        "nansum",
+        "nanprod",
+    ],
+)
+def test_reductions_toplevel(arr, func):
+    # var and std need __array_function__
+    result = getattr(da, func)(arr, axis=0)
+    expected = getattr(np, func)(arr.compute(), axis=0)
+    assert_eq(result, expected)
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index 6cd2d215c..1b7667638 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -120,24 +120,6 @@ def use_kwargs(a, kwarg=None):
     return a + kwarg
 
 
-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'))>"
-
-    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(...)>"
-
-    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 test_convert_legacy_dsk():
     def func(*args):
         return "-".join(args)
@@ -265,16 +247,87 @@ class SerializeOnlyOnce:
 
 def test_pickle():
 
+    def assert_slots_equal(a, b):
+        def get_all_slots(obj):
+            slots = set()
+            for cls in obj.__class__.mro():
+                slots.update(getattr(cls, "__slots__", ()))
+            return slots
+
+        all_slots = get_all_slots(a) | get_all_slots(a)
+        assert all_slots == get_all_slots(a) == get_all_slots(a)
+        assert all(getattr(a, slot) == getattr(b, slot) for slot in all_slots)
+        assert not hasattr(a, "__dict__")
+        assert not hasattr(b, "__dict__")
+
     t1 = Task("key-1", func, "a", "b")
     t2 = Task("key-2", func, "c", "d")
 
     rtt1 = pickle.loads(pickle.dumps(t1))
+    assert repr(rtt1) == repr(t1)
     rtt2 = pickle.loads(pickle.dumps(t2))
+
+    assert_slots_equal(t1, rtt1)
+    assert_slots_equal(t2, rtt2)
     assert t1 == rtt1
     assert t1.func == rtt1.func
     assert t1.func is rtt1.func
     assert t1.func is rtt2.func
 
+    l = Tuple(t1, t2)
+    rtl = pickle.loads(pickle.dumps(l))
+    assert l == rtl
+    assert l() == rtl()
+    d = Dict(key=t1)
+    rtd = pickle.loads(pickle.dumps(d))
+    assert d == rtd
+    assert d() == rtd()
+
+
+def test_pickle_size():
+    # We will serialize many of these objects which drives both memory usage and
+    # serialization runtime performance.
+    # Reducing pickle size is beneficial but the numbers below are determined
+    # empirically
+    # Analyzing the output with pickletools.dis is useful to debug memoization
+    # and serialization by value
+
+    a = Alias("a", "b")
+    # We cannot shrink it to nothing
+    assert len(pickle.dumps(a)) < 55
+    b = Alias("b", "c")
+    # But most of it should be overhead that is memoized
+    assert len(pickle.dumps((a, b))) <= 70
+
+    # Pickle should be able to memoize this. On py3.10 that's 2 additional bytes
+    assert len(pickle.dumps((a, b, b))) <= len(pickle.dumps((a, b))) + 10
+
+    t1 = Task("key-1", func, "a", "b")
+    assert len(pickle.dumps(t1)) < 120
+
+    t2 = Task("key-2", func, TaskRef("key-1"), "c")
+    assert len(pickle.dumps(t2)) < 140
+
+    assert len(pickle.dumps((t1, t2))) < 170
+
+    l = List(t1, t2)
+    assert len(pickle.dumps(l)) <= 272
+
+    sizes = []
+    growth = []
+    inner = List(t1, t2)
+    for depth in range(20):
+        inner = List(inner, t1)
+        size = len(pickle.dumps(inner))
+        if len(sizes) > 0:
+            growth.append(size - sizes[-1][1])
+        sizes.append((depth, size))
+    growth = set(growth)
+    # If this breaks, something cannot be memoized. That's very concerning
+    assert len(growth) == 1
+    # If this goes up, that's not great but not a disaster
+    assert growth.pop() <= 32
+
 
 def test_tokenize():
     t = Task("key-1", func, "a", "b")