Repository: dask/dask
Pull request #11692 (2025-01-23 11:55:07)

Implement unify_chunks and Rechunk
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`

sits on top of #11689

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 5m 33s ⏱️ - 18m 5s
 17 243 tests +  102   15 975 ✅ +    14   1 268 💤 +    88  0 ❌ ±0 
208 945 runs   - 2 267  178 448 ✅  - 15 741  30 497 💤 +13 474  0 ❌ ±0 

Results for commit 6c822007. ± Comparison against base commit d8eb38e9.

[test-results]:data:application/gzip;base64,H4sIAE03kmcC/03MTQ6DIBCG4asY1l0gMgz0Mg2MmJD60yCsmt69aC26fJ/JfG82hNGv7N60cGvYmkM6q8/RprDMW0vsukLlmvY7Cln7sWai/ckgnPgMrw2F0tUGG8ZivIKPcYmHxDxv04JrI+HoOo1aSl3xN91xafBvl+m9r8u0TFNIJZgiLQTn2LfeKEBpwXkYNAiwzhtyreJOEBn2+QLosJraGgEAAA==


Patch:
diff --git a/dask/_expr.py b/dask/_expr.py
index 9265bc1b6..6170691ee 100644
--- a/dask/_expr.py
+++ b/dask/_expr.py
@@ -40,9 +40,11 @@ def _unpack_collections(o):
 
 
 class Expr:
-    _parameters = []  # type: ignore
+    _parameters: list[str] = []
     _defaults: dict[str, Any] = {}
-    _instances = weakref.WeakValueDictionary()  # type: ignore
+    _instances: weakref.WeakValueDictionary[str, Expr] = weakref.WeakValueDictionary()
+
+    operands: list
 
     def __new__(cls, *args, **kwargs):
         operands = list(args)
@@ -99,16 +101,21 @@ class Expr:
         return os.linesep.join(self._tree_repr_lines())
 
     def analyze(self, filename: str | None = None, format: str | None = None) -> None:
+        from dask.dataframe.dask_expr._expr import Expr as DFExpr
         from dask.dataframe.dask_expr.diagnostics import analyze
 
-        return analyze(self, filename=filename, format=format)  # type: ignore
+        if not isinstance(self, DFExpr):
+            raise TypeError(
+                "analyze is only supported for dask.dataframe.Expr objects."
+            )
+        return analyze(self, filename=filename, format=format)
 
     def explain(
         self, stage: OptimizerStage = "fused", format: str | None = None
     ) -> None:
         from dask.dataframe.dask_expr.diagnostics import explain
 
-        return explain(self, stage, format)  # type: ignore
+        return explain(self, stage, format)
 
     def pprint(self):
         for line in self._tree_repr_lines():
@@ -266,7 +273,7 @@ class Expr:
             # Rewrite all of the children
             new_operands = []
             changed = False
-            for operand in expr.operands:  # type: ignore
+            for operand in expr.operands:
                 if isinstance(operand, Expr):
                     new = operand.rewrite(kind=kind, rewritten=rewritten)
                     rewritten[operand._name] = new
@@ -333,7 +340,7 @@ class Expr:
             # Rewrite all of the children
             new_operands = []
             changed = False
-            for operand in expr.operands:  # type: ignore
+            for operand in expr.operands:
                 if isinstance(operand, Expr):
                     # Bandaid for now, waiting for Singleton
                     dependents[operand._name].append(weakref.ref(expr))
@@ -354,6 +361,14 @@ class Expr:
 
         return expr
 
+    def optimize(self, fuse: bool = False) -> Expr:
+        stage: OptimizerStage = "fused" if fuse else "simplified-physical"
+
+        return optimize_until(self, stage)
+
+    def fuse(self) -> Expr:
+        return self
+
     def simplify(self) -> Expr:
         expr = self
         seen = set()
@@ -396,7 +411,7 @@ class Expr:
         # Lower all children
         new_operands = []
         changed = False
-        for operand in out.operands:  # type: ignore
+        for operand in out.operands:
             if isinstance(operand, Expr):
                 new = operand.lower_once(lowered)
                 if new._name != operand._name:
@@ -442,11 +457,11 @@ class Expr:
         return
 
     @functools.cached_property
-    def _funcname(self):
+    def _funcname(self) -> str:
         return funcname(type(self)).lower()
 
     @functools.cached_property
-    def _name(self):
+    def _name(self) -> str:
         return self._funcname + "-" + _tokenize_deterministic(*self.operands)
 
     @property
@@ -563,7 +578,7 @@ class Expr:
 
         changed = False
         new_operands = []
-        for i, operand in enumerate(self.operands):  # type: ignore
+        for i, operand in enumerate(self.operands):
             if i < len(self._parameters) and self._parameters[i] in substitutions:
                 new_operands.append(substitutions[self._parameters[i]])
                 changed = True
@@ -733,3 +748,61 @@ def collect_dependents(expr) -> defaultdict:
             stack.append(dep)
             dependents[dep._name].append(weakref.ref(node))
     return dependents
+
+
+def optimize(expr: Expr, fuse: bool = True) -> Expr:
+    """High level query optimization
+
+    This leverages three optimization passes:
+
+    1.  Class based simplification using the ``_simplify`` function and methods
+    2.  Blockwise fusion
+
+    Parameters
+    ----------
+    expr:
+        Input expression to optimize
+    fuse:
+        whether or not to turn on blockwise fusion
+
+    See Also
+    --------
+    simplify
+    optimize_blockwise_fusion
+    """
+    stage: OptimizerStage = "fused" if fuse else "simplified-physical"
+
+    return optimize_until(expr, stage)
+
+
+def optimize_until(expr: Expr, stage: OptimizerStage) -> Expr:
+    result = expr
+    if stage == "logical":
+        return result
+
+    # Simplify
+    expr = result.simplify()
+    if stage == "simplified-logical":
+        return expr
+
+    # Manipulate Expression to make it more efficient
+    expr = expr.rewrite(kind="tune", rewritten={})
+    if stage == "tuned-logical":
+        return expr
+
+    # Lower
+    expr = expr.lower_completely()
+    if stage == "physical":
+        return expr
+
+    # Simplify again
+    expr = expr.simplify()
+    if stage == "simplified-physical":
+        return expr
+
+    # Final graph-specific optimizations
+    expr = expr.fuse()
+    if stage == "fused":
+        return expr
+
+    raise ValueError(f"Stage {stage!r} not supported.")
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 80df3cfbd..5506e5553 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -224,7 +224,7 @@ def convert_legacy_task(
             new_args.append(new)
         return Task(key, func, *new_args)
     try:
-        if isinstance(task, (bytes, int, float, str, tuple)):
+        if isinstance(task, (int, float, str, tuple)):
             if task in all_keys:
                 if key is None:
                     return Alias(task)
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/creation.py b/dask/array/creation.py
index 0f69d165c..b906bc405 100644
--- a/dask/array/creation.py
+++ b/dask/array/creation.py
@@ -452,7 +452,7 @@ def arange(
         blockstop = start + ((elem_count + bs) * step)
         task = Task(
             (name, i),
-            partial(chunk.arange, like=like),
+            partial(chunk.arange, like=meta),
             blockstart,
             blockstop,
             step,
@@ -968,7 +968,7 @@ def expand_pad_value(array, pad_value):
     return pad_value
 
 
-def get_pad_shapes_chunks(array, pad_width, axes):
+def get_pad_shapes_chunks(array, pad_width, axes, mode):
     """
     Helper function for finding shapes and chunks of end pads.
     """
@@ -979,7 +979,12 @@ def get_pad_shapes_chunks(array, pad_width, axes):
     for d in axes:
         for i in range(2):
             pad_shapes[i][d] = pad_width[d][i]
-            pad_chunks[i][d] = (pad_width[d][i],)
+            if mode != "constant" or pad_width[d][i] == 0:
+                pad_chunks[i][d] = (pad_width[d][i],)
+            else:
+                pad_chunks[i][d] = normalize_chunks(
+                    (max(pad_chunks[i][d]),), (pad_width[d][i],)
+                )[0]
 
     pad_shapes = [tuple(s) for s in pad_shapes]
     pad_chunks = [tuple(c) for c in pad_chunks]
@@ -1021,7 +1026,9 @@ def pad_edge(array, pad_width, mode, **kwargs):
 
     result = array
     for d in range(array.ndim):
-        pad_shapes, pad_chunks = get_pad_shapes_chunks(result, pad_width, (d,))
+        pad_shapes, pad_chunks = get_pad_shapes_chunks(
+            result, pad_width, (d,), mode=mode
+        )
         pad_arrays = [result, result]
 
         if mode == "constant":
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:
diff --git a/dask/core.py b/dask/core.py
index bf0c2dceb..dd61cc980 100644
--- a/dask/core.py
+++ b/dask/core.py
@@ -149,7 +149,7 @@ def keys_in_tasks(keys: Collection[Key], tasks: Iterable[Any], as_list: bool = F
 def iskey(key: object) -> bool:
     """Return True if the given object is a potential dask key; False otherwise.
 
-    The definition of a key in a Dask graph is any str, bytes, int, float, or tuple
+    The definition of a key in a Dask graph is any str, int, float, or tuple
     thereof.
 
     See Also
@@ -161,7 +161,7 @@ def iskey(key: object) -> bool:
     typ = type(key)
     if typ is tuple:
         return all(iskey(i) for i in cast(tuple, key))
-    return typ in {bytes, int, float, str}
+    return typ in {int, float, str}
 
 
 def validate_key(key: object) -> None:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index a0aa0d9fe..32d8bbcb6 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -18,6 +18,7 @@ from dask.dataframe.utils import (
     has_known_categories,
     is_dataframe_like,
     is_index_like,
+    is_scalar,
     is_series_like,
     meta_frame_constructor,
     meta_series_constructor,
@@ -95,7 +96,13 @@ def apply_and_enforce(*args, **kwargs):
     func = kwargs.pop("_func")
     meta = kwargs.pop("_meta")
     df = func(*args, **kwargs)
-    if is_dataframe_like(df) or is_series_like(df) or is_index_like(df):
+
+    if any(
+        bool(is_dataframe_like(obj) or is_series_like(obj) or is_index_like(obj))
+        for obj in [df, meta]
+    ):
+        if is_scalar(df):
+            df = pd.Series(df)
         if not len(df):
             return meta
         if is_dataframe_like(df):
diff --git a/dask/dataframe/dask_expr/_collection.py b/dask/dataframe/dask_expr/_collection.py
index 70a18d8a3..4a2656bc5 100644
--- a/dask/dataframe/dask_expr/_collection.py
+++ b/dask/dataframe/dask_expr/_collection.py
@@ -106,7 +106,6 @@ from dask.dataframe.dask_expr._util import (
     _raise_if_object_series,
     _validate_axis,
     get_specified_shuffle,
-    is_scalar,
 )
 from dask.dataframe.dask_expr.io import FromPandasDivisions, FromScalars
 from dask.dataframe.dispatch import (
@@ -121,6 +120,7 @@ from dask.dataframe.utils import (
     has_known_categories,
     index_summary,
     insert_meta_param_description,
+    is_scalar,
     meta_frame_constructor,
     meta_series_constructor,
     pyarrow_strings_enabled,
@@ -397,6 +397,8 @@ class FrameBase(DaskMethodsMixin):
 
     def __getitem__(self, other):
         if isinstance(other, FrameBase):
+            if not expr.are_co_aligned(self.expr, other.expr):
+                return new_collection(expr.FilterAlign(self, other))
             return new_collection(self.expr.__getitem__(other.expr))
         elif isinstance(other, slice):
             from pandas.api.types import is_float_dtype
diff --git a/dask/dataframe/dask_expr/_expr.py b/dask/dataframe/dask_expr/_expr.py
index fe7390bdb..763f07a1b 100644
--- a/dask/dataframe/dask_expr/_expr.py
+++ b/dask/dataframe/dask_expr/_expr.py
@@ -37,7 +37,6 @@ from dask.dataframe.dask_expr._util import (
     _calc_maybe_new_divisions,
     _convert_to_list,
     _tokenize_partial,
-    is_scalar,
 )
 from dask.dataframe.dispatch import make_meta, meta_nonempty
 from dask.dataframe.rolling import CombinedOutput, _head_timedelta, overlap_chunk
@@ -45,6 +44,7 @@ from dask.dataframe.shuffle import drop_overlap, get_overlap
 from dask.dataframe.utils import (
     clear_known_categories,
     drop_by_shallow_copy,
+    is_scalar,
     raise_on_meta_error,
     valid_divisions,
 )
@@ -61,6 +61,8 @@ from dask.utils import (
     random_state_data,
 )
 
+optimize = core.optimize
+
 
 class Expr(core.Expr):
     """Primary class for all Expressions
@@ -176,8 +178,6 @@ class Expr(core.Expr):
 
     def __getitem__(self, other):
         if isinstance(other, Expr):
-            if not are_co_aligned(self, other):
-                return FilterAlign(self, other)
             return Filter(self, other)
         else:
             return Projection(self, other)  # df[["a", "b", "c"]]
@@ -520,7 +520,12 @@ class Expr(core.Expr):
     def _filter_simplification(self, parent, predicate=None):
         if predicate is None:
             predicate = parent.predicate.substitute(self, self.frame)
-        return type(self)(self.frame[predicate], *self.operands[1:])
+        if are_co_aligned(self.frame, predicate):
+            # Only do this if we are aligned
+            return type(self)(self.frame[predicate], *self.operands[1:])
+
+    def fuse(self):
+        return optimize_blockwise_fusion(self)
 
 
 class Literal(Expr):
@@ -1626,14 +1631,26 @@ class Where(Elemwise):
 
 
 def _check_divisions(df, i, division_min, division_max, last):
+    if not len(df):
+        return df
+    if is_index_like(df):
+        index = df
+    else:
+        try:
+            index = df.index.get_level_values(0)
+        except AttributeError:
+            index = df.index
     # Check divisions
-    real_min = df.index.min()
-    real_max = df.index.max()
+    real_min = index.min()
+    real_max = index.max()
     # Upper division of the last partition is often set to
     # the max value. For all other partitions, the upper
     # division should be greater than the maximum value.
-    valid_min = real_min >= division_min
-    valid_max = (real_max <= division_max) if last else (real_max < division_max)
+    valid_min = valid_max = True
+    if not pd.isna(division_min):
+        valid_min = real_min >= division_min
+    if not pd.isna(division_max):
+        valid_max = (real_max <= division_max) if last else (real_max < division_max)
     if not (valid_min and valid_max):
         raise RuntimeError(
             f"`enforce_runtime_divisions` failed for partition {i}."
@@ -1660,7 +1677,7 @@ class EnforceRuntimeDivisions(Blockwise):
             self.divisions[index + 1],
             index == (self.npartitions - 1),
         ]
-        return (self.operation,) + tuple(args)  # type: ignore
+        return Task(name, self.operation, *args)
 
 
 class Abs(Elemwise):
@@ -3074,7 +3091,7 @@ class _DelayedExpr(Expr):
 
 
 class DelayedsExpr(Expr):
-    _parameters = []  # type: ignore
+    _parameters = []
 
     def __init__(self, *delayed_objects):
         self.operands = delayed_objects
@@ -3108,64 +3125,6 @@ def normalize_expression(expr):
     return expr._name
 
 
-def optimize_until(expr: Expr, stage: core.OptimizerStage) -> Expr:
-    result = expr
-    if stage == "logical":
-        return result
-
-    # Simplify
-    expr = result.simplify()  # type: ignore
-    if stage == "simplified-logical":
-        return expr
-
-    # Manipulate Expression to make it more efficient
-    expr = expr.rewrite(kind="tune", rewritten={})
-    if stage == "tuned-logical":
-        return expr
-
-    # Lower
-    expr = expr.lower_completely()  # type: ignore
-    if stage == "physical":
-        return expr
-
-    # Simplify again
-    expr = expr.simplify()  # type: ignore
-    if stage == "simplified-physical":
-        return expr
-
-    # Final graph-specific optimizations
-    expr = optimize_blockwise_fusion(expr)
-    if stage == "fused":
-        return expr
-
-    raise ValueError(f"Stage {stage!r} not supported.")
-
-
-def optimize(expr: Expr, fuse: bool = True) -> Expr:
-    """High level query optimization
-
-    This leverages three optimization passes:
-
-    1.  Class based simplification using the ``_simplify`` function and methods
-    2.  Blockwise fusion
-
-    Parameters
-    ----------
-    expr:
-        Input expression to optimize
-    fuse:
-        whether or not to turn on blockwise fusion
-
-    See Also
-    --------
-    simplify
-    optimize_blockwise_fusion
-    """
-    stage: core.OptimizerStage = "fused" if fuse else "simplified-physical"
-
-    return optimize_until(expr, stage)
-
-
 def is_broadcastable(dfs, s):
     """
     This Series is broadcastable against another dataframe in the sequence
@@ -3920,7 +3879,9 @@ def plain_column_projection(expr, parent, dependents, additional_columns=None):
         # we are accesing the index
         column_union = []
 
-    if column_union == expr.frame.columns:
+    if column_union == expr.frame.columns or not column_union and expr.ndim < 2:
+        # this projection is for the index, but the elements are unknown, so
+        # don't project
         return
     result = type(expr)(expr.frame[column_union], *expr.operands[1:])
     if column_union == parent.operand("columns"):
diff --git a/dask/dataframe/dask_expr/_groupby.py b/dask/dataframe/dask_expr/_groupby.py
index d56cacd39..8c555c0e8 100644
--- a/dask/dataframe/dask_expr/_groupby.py
+++ b/dask/dataframe/dask_expr/_groupby.py
@@ -45,7 +45,6 @@ from dask.dataframe.dask_expr._util import (
     PANDAS_GE_300,
     _convert_to_list,
     get_specified_shuffle,
-    is_scalar,
 )
 from dask.dataframe.dispatch import concat, make_meta, meta_nonempty
 from dask.dataframe.groupby import (
@@ -82,7 +81,7 @@ from dask.dataframe.groupby import (
     _var_agg,
     _var_chunk,
 )
-from dask.dataframe.utils import insert_meta_param_description
+from dask.dataframe.utils import insert_meta_param_description, is_scalar
 from dask.typing import Key
 from dask.utils import M, derived_from, is_index_like
 
diff --git a/dask/dataframe/dask_expr/_indexing.py b/dask/dataframe/dask_expr/_indexing.py
index 12bd43f04..7758c032d 100644
--- a/dask/dataframe/dask_expr/_indexing.py
+++ b/dask/dataframe/dask_expr/_indexing.py
@@ -20,7 +20,6 @@ from dask.dataframe.dask_expr._expr import (
     Projection,
     are_co_aligned,
 )
-from dask.dataframe.dask_expr._util import is_scalar
 from dask.dataframe.dispatch import meta_nonempty
 from dask.dataframe.indexing import (
     _coerce_loc_index,
@@ -28,6 +27,7 @@ from dask.dataframe.indexing import (
     _partition_of_index_value,
     _partitions_of_index_values,
 )
+from dask.dataframe.utils import is_scalar
 from dask.typing import Key
 from dask.utils import is_arraylike, is_series_like
 
diff --git a/dask/dataframe/dask_expr/_merge.py b/dask/dataframe/dask_expr/_merge.py
index 190cde1e2..d6fecac69 100644
--- a/dask/dataframe/dask_expr/_merge.py
+++ b/dask/dataframe/dask_expr/_merge.py
@@ -33,7 +33,7 @@ from dask.dataframe.dask_expr._shuffle import (
     _is_numeric_cast_type,
     _select_columns_or_index,
 )
-from dask.dataframe.dask_expr._util import _convert_to_list, is_scalar
+from dask.dataframe.dask_expr._util import _convert_to_list
 from dask.dataframe.dispatch import make_meta, meta_nonempty
 from dask.dataframe.multi import (
     _concat_wrapper,
@@ -42,6 +42,7 @@ from dask.dataframe.multi import (
     merge_chunk,
 )
 from dask.dataframe.shuffle import partitioning_index
+from dask.dataframe.utils import is_scalar
 from dask.tokenize import _tokenize_deterministic
 from dask.typing import Key
 from dask.utils import apply, get_default_shuffle_method
diff --git a/dask/dataframe/dask_expr/_reductions.py b/dask/dataframe/dask_expr/_reductions.py
index 5ba667896..da9fb6a2a 100644
--- a/dask/dataframe/dask_expr/_reductions.py
+++ b/dask/dataframe/dask_expr/_reductions.py
@@ -38,8 +38,8 @@ from dask.dataframe.dask_expr._expr import (
     determine_column_projection,
     plain_column_projection,
 )
-from dask.dataframe.dask_expr._util import is_scalar
 from dask.dataframe.dispatch import make_meta, meta_nonempty
+from dask.dataframe.utils import is_scalar
 from dask.tokenize import _tokenize_deterministic
 from dask.typing import no_default
 from dask.utils import M, apply, funcname
diff --git a/dask/dataframe/dask_expr/_repartition.py b/dask/dataframe/dask_expr/_repartition.py
index fe13dd362..d5a9dd5d7 100644
--- a/dask/dataframe/dask_expr/_repartition.py
+++ b/dask/dataframe/dask_expr/_repartition.py
@@ -149,6 +149,10 @@ class Repartition(Expr):
         if isinstance(parent, Filter) and self._filter_passthrough_available(
             parent, dependents
         ):
+            if self._name == parent.predicate._name:
+                # We shouldn't push through the predicate, these pushdowns should
+                # always come from frame.
+                return
             return self._filter_simplification(parent)
         if isinstance(parent, Projection):
             return plain_column_projection(self, parent, dependents)
diff --git a/dask/dataframe/dask_expr/_util.py b/dask/dataframe/dask_expr/_util.py
index aef4d8e90..c9de28dab 100644
--- a/dask/dataframe/dask_expr/_util.py
+++ b/dask/dataframe/dask_expr/_util.py
@@ -2,10 +2,9 @@ from __future__ import annotations
 
 import functools
 from collections import OrderedDict, UserDict
-from collections.abc import Hashable, Iterable, Sequence
+from collections.abc import Hashable
 from typing import Any, Literal, TypeVar, cast
 
-import numpy as np
 import pandas as pd
 from packaging.version import Version
 
@@ -79,22 +78,6 @@ def _convert_to_list(column) -> list | None:
     return column
 
 
-def is_scalar(x):
-    # np.isscalar does not work for some pandas scalars, for example pd.NA
-    if isinstance(x, (Sequence, Iterable)) and not isinstance(x, str):
-        return False
-    elif hasattr(x, "dtype"):
-        return isinstance(x, np.ScalarType)
-    if isinstance(x, dict):
-        return False
-    if isinstance(x, (str, int)) or x is None:
-        return True
-
-    from dask.dataframe.dask_expr._expr import Expr
-
-    return not isinstance(x, Expr)
-
-
 def _tokenize_partial(expr, ignore: list | None = None) -> str:
     # Helper function to "tokenize" the operands
     # that are not in the `ignore` list
diff --git a/dask/dataframe/dask_expr/diagnostics/_analyze.py b/dask/dataframe/dask_expr/diagnostics/_analyze.py
index 29d2cadc4..0dfd822ca 100644
--- a/dask/dataframe/dask_expr/diagnostics/_analyze.py
+++ b/dask/dataframe/dask_expr/diagnostics/_analyze.py
@@ -7,12 +7,12 @@ import pandas as pd
 
 from dask.base import DaskMethodsMixin
 from dask.dataframe.dask_expr._expr import Blockwise, Expr
-from dask.dataframe.dask_expr._util import is_scalar
 from dask.dataframe.dask_expr.diagnostics._explain import (
     _add_graphviz_edges,
     _explain_info,
 )
 from dask.dataframe.dask_expr.io.io import FusedIO
+from dask.dataframe.utils import is_scalar
 from dask.sizeof import sizeof
 from dask.tokenize import _tokenize_deterministic
 from dask.utils import format_bytes, import_required
diff --git a/dask/dataframe/dask_expr/diagnostics/_explain.py b/dask/dataframe/dask_expr/diagnostics/_explain.py
index 1d5beebc3..5912b82fe 100644
--- a/dask/dataframe/dask_expr/diagnostics/_explain.py
+++ b/dask/dataframe/dask_expr/diagnostics/_explain.py
@@ -1,7 +1,8 @@
 from __future__ import annotations
 
-from dask._expr import OptimizerStage
-from dask.dataframe.dask_expr._expr import Expr, Projection, optimize_until
+from dask._expr import Expr as BaseExpr
+from dask._expr import OptimizerStage, optimize_until
+from dask.dataframe.dask_expr._expr import Expr, Projection
 from dask.dataframe.dask_expr._merge import Merge
 from dask.dataframe.dask_expr.io.parquet import ReadParquet
 from dask.utils import funcname, import_required
@@ -17,7 +18,7 @@ STAGE_LABELS: dict[OptimizerStage, str] = {
 
 
 def explain(
-    expr: Expr, stage: OptimizerStage = "fused", format: str | None = None
+    expr: BaseExpr, stage: OptimizerStage = "fused", format: str | None = None
 ) -> None:
     graphviz = import_required(
         "graphviz", "graphviz is a required dependency for using the explain method."
@@ -34,7 +35,7 @@ def explain(
     expr = optimize_until(expr, stage)
 
     seen = set(expr._name)
-    stack = [expr]
+    stack: list[BaseExpr] = [expr]
 
     while stack:
         node = stack.pop()
@@ -73,7 +74,7 @@ def _add_graphviz_edges(explain_info, graph):
         graph.edge(dep, name)
 
 
-def _explain_info(expr: Expr):
+def _explain_info(expr: BaseExpr):
     return {
         "name": expr._name,
         "label": funcname(type(expr)),
@@ -82,7 +83,10 @@ def _explain_info(expr: Expr):
     }
 
 
-def _explain_details(expr: Expr):
+def _explain_details(expr: BaseExpr):
+
+    if not isinstance(expr, Expr):
+        return {}
     details = {"npartitions": expr.npartitions}
 
     if isinstance(expr, Merge):
@@ -96,10 +100,10 @@ def _explain_details(expr: Expr):
     return details
 
 
-def _explain_dependencies(expr: Expr) -> list[tuple[str, str]]:
+def _explain_dependencies(expr: BaseExpr) -> list[tuple[str, str]]:
     dependencies = []
     for i, operand in enumerate(expr.operands):
-        if not isinstance(operand, Expr):
+        if not isinstance(operand, BaseExpr):
             continue
         param = expr._parameters[i] if i < len(expr._parameters) else ""
         dependencies.append((str(param), operand._name))
diff --git a/dask/dataframe/dask_expr/io/io.py b/dask/dataframe/dask_expr/io/io.py
index e0edf672d..ab416156f 100644
--- a/dask/dataframe/dask_expr/io/io.py
+++ b/dask/dataframe/dask_expr/io/io.py
@@ -445,6 +445,8 @@ class FromPandas(PartitionsFiltered, BlockwiseIO):
             try:
                 return list(self.frame.columns)
             except AttributeError:
+                if self.ndim == 1:
+                    return [self.name]
                 return []
         else:
             return _convert_to_list(columns_operand)
diff --git a/dask/dataframe/utils.py b/dask/dataframe/utils.py
index 1911e5de7..7260ae0d6 100644
--- a/dask/dataframe/utils.py
+++ b/dask/dataframe/utils.py
@@ -5,7 +5,7 @@ import re
 import sys
 import textwrap
 import traceback
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from contextlib import contextmanager
 from numbers import Number
 from typing import TypeVar, overload
@@ -15,7 +15,7 @@ import pandas as pd
 from pandas.api.types import is_dtype_equal
 
 import dask
-from dask.base import get_scheduler, is_dask_collection
+from dask.base import is_dask_collection
 from dask.core import get_deps
 from dask.dataframe._compat import tm  # noqa: F401
 from dask.dataframe.dispatch import (  # noqa : F401
@@ -43,6 +43,22 @@ except ImportError:
     pass
 
 
+def is_scalar(x):
+    # np.isscalar does not work for some pandas scalars, for example pd.NA
+    if isinstance(x, (Sequence, Iterable)) and not isinstance(x, str):
+        return False
+    elif hasattr(x, "dtype"):
+        return isinstance(x, np.ScalarType)
+    if isinstance(x, dict):
+        return False
+    if isinstance(x, (str, int)) or x is None:
+        return True
+
+    from dask.dataframe.dask_expr._expr import Expr
+
+    return not isinstance(x, Expr)
+
+
 def is_integer_na_dtype(t):
     dtype = getattr(t, "dtype", t)
     types = (
@@ -587,24 +603,7 @@ def assert_divisions(ddf, scheduler=None):
     if not getattr(ddf, "known_divisions", False):
         return
 
-    def index(x):
-        if is_index_like(x):
-            return x
-        try:
-            return x.index.get_level_values(0)
-        except AttributeError:
-            return x.index
-
-    get = get_scheduler(scheduler=scheduler, collections=[type(ddf)])
-    results = get(ddf.dask, ddf.__dask_keys__())
-    for i, df in enumerate(results[:-1]):
-        if len(df):
-            assert index(df).min() >= ddf.divisions[i]
-            assert index(df).max() < ddf.divisions[i + 1]
-
-    if len(results[-1]):
-        assert index(results[-1]).min() >= ddf.divisions[-2]
-        assert index(results[-1]).max() <= ddf.divisions[-1]
+    ddf.enforce_runtime_divisions().compute(scheduler=scheduler)
 
 
 def assert_sane_keynames(ddf):
diff --git a/dask/typing.py b/dask/typing.py
index 6c67ff746..8a970f87e 100644
--- a/dask/typing.py
+++ b/dask/typing.py
@@ -25,7 +25,7 @@ CollType_co = TypeVar("CollType_co", bound="DaskCollection", covariant=True)
 PostComputeCallable = Callable
 
 
-Key: TypeAlias = Union[str, bytes, int, float, tuple["Key", ...]]
+Key: TypeAlias = Union[str, int, float, tuple["Key", ...]]
 # FIXME: This type is a little misleading. Low level graphs are often
 # MutableMappings but HLGs are not
 Graph: TypeAlias = Mapping[Key, Any]
diff --git a/docs/source/array.rst b/docs/source/array.rst
index 90d5524fe..26fc7a316 100644
--- a/docs/source/array.rst
+++ b/docs/source/array.rst
@@ -38,8 +38,18 @@ these blocked algorithms using Dask graphs.
 Examples
 --------
 
-Visit https://examples.dask.org/array.html to see and run examples using
-Dask Array.
+Dask Array is used across a wide variety of applications — anywhere where working
+with large array dataset.
+
+This
+`Analyzing the National Water Model with Xarray, Dask, and Coiled <https://docs.coiled.io/user_guide/xarray.html?utm_source=dask-docs&utm_medium=array>`_
+example process 6 TB of geospatial data on a cluster using Xarray and Dask Array.
+The cluster in this example is deployed with
+`Coiled <https://coiled.io/?utm_source=dask-docs&utm_medium=array>`_,
+but there are many options for managing and deploying Dask.
+See our :doc:`deploying` documentation for more information on deployment options.
+
+You can also visit https://examples.dask.org/array.html for a collection of additional examples.
 
 Design
 ------
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index e84a163e0..92b1da12d 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -5,6 +5,69 @@ Changelog
 
     This is not exhaustive. For an exhaustive list of changes, see the git log.
 
+.. _v2025.2.0:
+
+2025.2.0
+--------
+
+Highlights
+^^^^^^^^^^
+This release includes a critical fix that fixes a deadlock that can arise
+when seceded task are rescheduled, or cancelled and resubmitted, e.g. due
+to a worker being lost.
+
+See :pr-distributed:`8991` by `Hendrik Makait`_ for more details.
+
+.. dropdown:: Additional changes
+
+  - Add big array example (:pr:`11744`) `James Bourbeau`_
+  - Fix exploding chunksizes in pad for constant padding (:pr:`11743`) `Patrick Hoefler`_
+  - Move optimize method to base class (:pr:`11742`) `Florian Jetter`_
+  - Add changelog entry for fixed deadlock (:pr:`11741`) `Hendrik Makait`_
+  - Fix graph creation in ``dask-expr`` ``to_delayed`` (:pr:`11739`) `Patrick Hoefler`_
+  - Remove culling from delayed optimisation (:pr:`11737`) `Patrick Hoefler`_
+  - Compute meta for from_map on the cluster (:pr:`11738`) `Patrick Hoefler`_
+  - Bugs in ``__setitem__`` with dask bool mask (:pr:`11728`) `Guido Imperiale`_
+  - Implement infrastructure, random, blockwise and Elemwise  (:pr:`11689`) `Patrick Hoefler`_
+  - ``array`` / ``asarray`` with both ``like=`` and ``dtype=`` (:pr:`11733`) `Guido Imperiale`_
+  - Fix annotations warnings test (:pr:`11734`) `Patrick Hoefler`_
+  - Catch warnings when writing to remote storage with to_parquet (:pr:`11731`) `Patrick Hoefler`_
+  - Remove LocalCluster from tests (:pr:`11729`) `Patrick Hoefler`_
+  - Fix partition pruning when using from_array (:pr:`11725`) `Patrick Hoefler`_
+  - Fix concatentation with mixed dtype columns (:pr:`11727`) `Patrick Hoefler`_
+  - ``arange``: fix extreme values (:pr:`11707`) `Guido Imperiale`_
+  - Graph corruption on scalar ``getitem`` -> ``setitem`` (:pr:`11723`) `Guido Imperiale`_
+  - Never share buffers after compute() (:pr:`11697`) `Guido Imperiale`_
+  - Extract Dask Array from xarray DataArray in from_array (:pr:`11712`) `Patrick Hoefler`_
+  - ``arange``: support kwargs (:pr:`11710`) `Guido Imperiale`_
+  - Ensure ``normalize_token`` is threadsafe (:pr:`11709`) `Florian Jetter`_
+  - Expand advise for instance types and processes (:pr:`11705`) `Florian Jetter`_
+  - Drop legacy timeseries implementation (:pr:`11704`) `Florian Jetter`_
+  - Update Dask Cloud Provider documentation to include Nebius as a supported cloud option (:pr:`11703`) `Alexander`_
+  - Fix ``normalize_chunks`` when squashing into a single chunk (:pr:`11702`) `Patrick Hoefler`_
+  - Fix positional indexing with ``newaxis`` (:pr:`11699`) `Patrick Hoefler`_
+  - Set array backend in scipy-sparse-indexing (:pr:`11700`) `Tom Augspurger`_
+  - Fix ``value_counts`` shuffling strategy (:pr:`11698`) `Patrick Hoefler`_
+  - Disentangle core expression class from dataframe specific code (:pr:`11688`) `Patrick Hoefler`_
+  - Bump ``conda-incubator/setup-miniconda`` from 3.1.0 to 3.1.1 (:pr:`11685`)
+  - Fixup dataframe conversion from array methods (:pr:`11684`) `Patrick Hoefler`_
+  - Remove remaining artifacts of ``fastparquet`` (:pr:`11682`) `Patrick Hoefler`_
+
+  - Remove traceback from ``sizeof`` failure warning (:pr-distributed:`9006`) `Jacob Tomlinson`_
+  - Hotfix: Ignore negative occupancy (:pr-distributed:`9012`) `Hendrik Makait`_
+  - Remove expensive tokenization for key uniqueness check (:pr-distributed:`9009`) `Patrick Hoefler`_
+  - Fix CI for changes in ``from_map`` (:pr-distributed:`9011`) `Patrick Hoefler`_
+  - Avoid handling stale long-running messages on scheduler (:pr-distributed:`8991`) `Hendrik Makait`_
+  - Bump ``test_stress`` timeout (:pr-distributed:`9002`) `Tom Augspurger`_
+  - Poll in ``test_rmm_metrics`` test (:pr-distributed:`9004`) `Tom Augspurger`_
+  - Cache occupancy in ``WorkStealing.balance()`` (:pr-distributed:`9005`) `Hendrik Makait`_
+  - Homogeneous balancing by accounting for in-flight requests (:pr-distributed:`9003`) `Hendrik Makait`_
+  - Consistent estimation of task duration between stealing, adaptive and occupancy calculation (:pr-distributed:`9000`) `Hendrik Makait`_
+  - Increase default work-stealing interval by 10x (:pr-distributed:`8997`) `Hendrik Makait`_
+  - Remove occupancy plot from status dashboard (:pr-distributed:`8995`) `Hendrik Makait`_
+  - Bump ``conda-incubator/setup-miniconda`` from 3.1.0 to 3.1.1 (:pr-distributed:`8990`)
+
+
 .. _v2025.1.0:
 
 2025.1.0
@@ -9206,3 +9269,5 @@ Other
 .. _`Ilan Gold`: https://github.com/ilan-gold
 .. _`Jean-Baptiste Bayle`: https://github.com/j2bbayle
 .. _`dchudz`: https://github.com/dchudz
+.. _`Guido Imperiale`: https://github.com/crusaderky
+.. _`Alexander`: https://github.com/SalikovAlex
diff --git a/pyproject.toml b/pyproject.toml
index 1e8a62480..6548c4bc1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -55,7 +55,7 @@ dataframe = [
     "pandas >= 2.0",
     "pyarrow>=14.0.1",
 ]
-distributed = ["distributed == 2025.1.0"]
+distributed = ["distributed == 2025.2.0"]
 diagnostics = [
     "bokeh >= 3.1.0",
     "jinja2 >= 2.10.3",


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 958c6e13e..f086d4405 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -125,7 +125,8 @@ jobs:
         # ensure this runs even if pytest fails
         if: >
           always() &&
-          (steps.run_tests.outcome == 'success' || steps.run_tests.outcome == 'failure')
+          (steps.run_tests.outcome == 'success' || steps.run_tests.outcome == 'failure') &&
+          matrix.array-expr == 'false'
         uses: actions/upload-artifact@v4
         with:
           name: ${{ matrix.os }}-${{ matrix.environment }}-${{ matrix.extra }}
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/array/tests/test_creation.py b/dask/array/tests/test_creation.py
index fef3272e8..fba157f8a 100644
--- a/dask/array/tests/test_creation.py
+++ b/dask/array/tests/test_creation.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import pickle
+
 import pytest
 
 from dask._task_spec import Alias
@@ -991,6 +993,23 @@ def test_pad_udf(kwargs):
     assert_eq(np_r, da_r)
 
 
+def test_pad_constant_chunksizes():
+    array = dask.array.ones((10, 10), chunks=(1, 1))
+    result = dask.array.pad(
+        array, ((0, 16 - 10), (0, 0)), mode="constant", constant_values=0
+    )
+    assert tuple(map(max, result.chunks)) == (1, 1)
+    assert_eq(
+        result,
+        np.pad(
+            array.compute(),
+            mode="constant",
+            constant_values=0,
+            pad_width=((0, 16 - 10), (0, 0)),
+        ),
+    )
+
+
 def test_auto_chunks():
     with dask.config.set({"array.chunk-size": "50 MiB"}):
         x = da.ones((10000, 10000))
@@ -1076,3 +1095,18 @@ def test_nan_full_like(val, shape_chunks, dtype):
         da.full_like(y1, val, dtype=dtype),
         np.full_like(y2, val, dtype=dtype),
     )
+
+
+@pytest.mark.parametrize(
+    "func", [da.array, da.asarray, da.asanyarray, da.arange, da.tri]
+)
+def test_like_forgets_graph(func):
+    """Test that array creation functions with like=x do not
+    internally store the graph of x
+    """
+    x = da.arange(3).map_blocks(lambda x: x)
+    with pytest.raises(Exception, match="local object"):
+        pickle.dumps(x)
+
+    a = func(1, like=x)
+    pickle.dumps(a)
diff --git a/dask/dataframe/dask_expr/tests/test_collection.py b/dask/dataframe/dask_expr/tests/test_collection.py
index a4d2d15c4..6f241ff38 100644
--- a/dask/dataframe/dask_expr/tests/test_collection.py
+++ b/dask/dataframe/dask_expr/tests/test_collection.py
@@ -22,7 +22,6 @@ from dask.dataframe.dask_expr import (
     expr,
     from_dict,
     from_pandas,
-    is_scalar,
     isna,
     optimize,
     to_datetime,
@@ -40,7 +39,7 @@ from dask.dataframe.dask_expr._shuffle import Shuffle
 from dask.dataframe.dask_expr.datasets import timeseries
 from dask.dataframe.dask_expr.io import FromPandas
 from dask.dataframe.dask_expr.tests._util import _backend_library, assert_eq, xfail_gpu
-from dask.dataframe.utils import UNKNOWN_CATEGORIES, pyarrow_strings_enabled
+from dask.dataframe.utils import UNKNOWN_CATEGORIES, is_scalar, pyarrow_strings_enabled
 from dask.utils import M
 
 # Set DataFrame backend for this module
@@ -2785,3 +2784,27 @@ def test_to_backend_simplify():
         assert str(df2.expr) != str(df[["y"]].expr)
         df3 = df2.simplify()
         assert str(df3.expr) == str(df[["y"]].expr)
+
+
+def test_getitem_triggering_unnecessary_alignment():
+    df = pd.DataFrame(
+        {
+            "A": [1, 2, 3, 4, 5, 6, 7, 8, 9],
+            "B": [9, 8, 7, 6, 5, 4, 3, 2, 1],
+            "C": [True, False, True] * 3,
+        },
+        columns=list("ABC"),
+    )
+    ddf = from_pandas(df, 2)
+    result = ddf[ddf.C]
+    expr = result.expr.optimize()
+    assert expr.__dask_graph__()
+    assert_eq(result, df[df.C])
+
+
+def test_projection_on_series():
+    pdf = pd.DataFrame(data={"a": [1, 3, 2]}).a
+    df = from_pandas(pdf, npartitions=1)
+    result = df.replace(1, 5)
+    assert_eq(result, pdf.replace(1, 5))
+    assert_eq(result.fillna(0), pdf.replace(1, 5).fillna(0))
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index e9716e420..52e8e8d15 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -903,6 +903,7 @@ def test_map_partitions():
         d.map_partitions(lambda df: 1),
         pd.Series([1, 1, 1], dtype=np.int64),
         check_divisions=False,
+        check_index=False,
     )
 
 
@@ -5353,3 +5354,19 @@ def test_array_to_df_conversion():
     result = arr.map_blocks(foo, meta=pd.DataFrame(np.random.random((1, 10))))
     expected = pd.DataFrame(arr.compute())
     assert_eq(result, expected, check_index=False)
+
+
+@pytest.mark.parametrize("npartitions", [1, 5])
+def test_map_partitions_always_series(npartitions):
+    pdf = pd.DataFrame({"x": range(50)})
+    ddf = dd.from_pandas(pdf, npartitions)
+
+    def assert_series(x):
+        assert isinstance(x, pd.Series)
+        return x
+
+    res = ddf.x.map_partitions(M.min).map_partitions(assert_series).compute()
+    assert len(res) == ddf.npartitions
+    assert min(res) == min(pdf.x)
+
+    assert ddf.x.map_partitions(M.min).count().compute() == ddf.npartitions