Repository: dask/dask
Pull request #11761 (2025-02-19 12:05:21)

skip test_visualize_int_overflow on windows
This is causing a brokenpipe error on windows. I suspect this is related to graphviz but I don't think it's worth debugging for us


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_order.py b/dask/tests/test_order.py
index 1536a2b3f..c1711d537 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -8,6 +8,7 @@ from packaging.version import Version
 
 import dask
 from dask import delayed
+from dask._compatibility import WINDOWS
 from dask.base import collections_to_dsk, key_split, visualize_dsk
 from dask.core import get_deps
 from dask.order import _connecting_to_roots, diagnostics, ndependencies, order
@@ -254,6 +255,7 @@ def _rechunk_merge_graph():
     }
 
 
+@pytest.mark.skipif(WINDOWS, reason="Broken pipe on Windows")
 def test_visualize_int_overflow():
     pytest.importorskip("graphviz")
     pytest.importorskip("matplotlib")
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")