diff --git a/flox/aggregations.py b/flox/aggregations.py
index 06a112a3..5bd5a137 100644
--- a/flox/aggregations.py
+++ b/flox/aggregations.py
@@ -14,7 +14,7 @@
 
 from . import aggregate_flox, aggregate_npg, xrutils
 from . import xrdtypes as dtypes
-from .lib import dask_array_type, sparse_array_type
+from .lib import sparse_array_type
 from .multiarray import MultiArray
 from .xrutils import notnull
 
@@ -1031,7 +1031,7 @@ def _initialize_aggregation(
 
 
 def is_supported_aggregation(array, func: str, **kwargs) -> bool:
-    if isinstance(array, dask_array_type):
+    if xrutils.is_duck_dask_array(array):
         # need to check the underlying array type
         array = array._meta
 
diff --git a/flox/core.py b/flox/core.py
index 15fd8944..cf4a2323 100644
--- a/flox/core.py
+++ b/flox/core.py
@@ -32,7 +32,7 @@
     _factorize_multiple,
     factorize_,
 )
-from .lib import sparse_array_type
+from .lib import contains_standalone_dask_array, is_standalone_dask_array, sparse_array_type
 from .rechunk import rechunk_for_blockwise
 from .reindex import (
     ReindexArrayType,
@@ -1154,7 +1154,13 @@ def groupby_reduce(
         if kwargs["fill_value"] is None:
             kwargs["fill_value"] = agg.fill_value[agg.name]
 
-        from .dask import dask_groupby_agg
+        dask_inputs = tuple(arg for arg in (array, by_) if is_duck_dask_array(arg))
+        if contains_standalone_dask_array(*dask_inputs):
+            if not all(is_standalone_dask_array(arg) for arg in dask_inputs):
+                raise TypeError("Cannot mix dask_array.Array with other Dask-backed array types.")
+            from .dask_array import dask_groupby_agg
+        else:
+            from .dask import dask_groupby_agg
 
         partial_agg = partial(dask_groupby_agg, **kwargs)
 
diff --git a/flox/dask_array.py b/flox/dask_array.py
new file mode 100644
index 00000000..cab852fe
--- /dev/null
+++ b/flox/dask_array.py
@@ -0,0 +1,441 @@
+"""dask-array expression integration for flox.
+
+This module is intentionally narrow: it mirrors the existing ``flox.dask``
+groupby construction, but uses dask-array's expression primitives when the
+input is a standalone ``dask_array.Array``.
+"""
+
+from __future__ import annotations
+
+import itertools
+import operator
+from collections.abc import Sequence
+from functools import cached_property, partial
+from numbers import Integral
+from typing import TYPE_CHECKING, Any
+
+import dask
+import dask_array as da
+import numpy as np
+import pandas as pd
+import toolz as tlz
+from dask._task_spec import Alias, Task, TaskRef
+from dask_array._collection import Array
+from dask_array._expr import ArrayExpr
+from dask_array._new_collection import new_collection
+
+from .cohorts import find_group_cohorts
+from .core import _get_chunk_reduction, _reduce_blockwise
+from .dask import (
+    _aggregate,
+    _expand_dims,
+    _extract_result,
+    _grouped_combine,
+    _normalize_indexes,
+    _simple_combine,
+    reindex_intermediates,
+)
+from .lib import _is_arg_reduction, _is_first_last_reduction, identity
+from .reindex import ReindexStrategy
+from .xrutils import is_duck_dask_array
+
+if TYPE_CHECKING:
+    from .aggregations import Aggregation
+    from .core import T_Axes, T_Engine, T_Method
+
+
+def is_dask_array(x: Any) -> bool:
+    return isinstance(x, Array)
+
+
+def contains_dask_array(*args: Any) -> bool:
+    return any(is_dask_array(arg) for arg in args)
+
+
+def _passthrough_chunk(x, axis=None, keepdims=None):
+    return x
+
+
+class _DirectChunkManager:
+    @property
+    def array_api(self):
+        return da
+
+    def from_array(self, data, chunks, **kwargs):
+        return da.from_array(data, chunks=chunks, **kwargs)
+
+    def map_blocks(self, func, *args, **kwargs):
+        return da.map_blocks(func, *args, **kwargs)
+
+    def blockwise(self, func, out_ind, *args, **kwargs):
+        return da.blockwise(func, out_ind, *args, **kwargs)
+
+    def unify_chunks(self, *args, **kwargs):
+        return da.unify_chunks(*args, **kwargs)
+
+
+_DIRECT_CHUNKMANAGER = _DirectChunkManager()
+
+
+def get_chunkmanager(*args: Any):
+    chunked_args = tuple(arg for arg in args if is_dask_array(arg))
+    if not chunked_args:
+        return _DIRECT_CHUNKMANAGER
+
+    try:
+        from xarray.namedarray.parallelcompat import get_chunked_array_type
+    except ImportError:
+        return _DIRECT_CHUNKMANAGER
+
+    try:
+        return get_chunked_array_type(*chunked_args)
+    except TypeError:
+        return _DIRECT_CHUNKMANAGER
+
+
+class FloxExtractGroups(ArrayExpr):
+    _parameters = ["array", "key", "_dtype"]
+
+    @cached_property
+    def dtype(self):
+        return np.dtype(self.operand("_dtype"))
+
+    @cached_property
+    def chunks(self):
+        return ((np.nan,),)
+
+    @cached_property
+    def _meta(self):
+        return np.array([], dtype=self.dtype)
+
+    def _layer(self):
+        first_block = self.array.ndim * (0,)
+        in_key = (self.array.name, *first_block)
+        out_key = (self._name, 0)
+        return {out_key: Task(out_key, operator.getitem, TaskRef(in_key), self.key)}
+
+
+class FloxCollapseBlocks(ArrayExpr):
+    _parameters = ["array", "axis", "group_chunks"]
+
+    @cached_property
+    def dtype(self):
+        return self.array.dtype
+
+    @cached_property
+    def chunks(self):
+        axis = self.axis
+        return self.array.chunks[: -len(axis)] + ((1,) * (len(axis) - 1),) + self.group_chunks
+
+    @cached_property
+    def _meta(self):
+        return self.array._meta
+
+    def _layer(self):
+        axis = self.axis
+        nblocks = tuple(self.array.numblocks[ax] for ax in axis)
+        layer = {}
+        for out_block in itertools.product(*(range(len(chunks)) for chunks in self.chunks)):
+            in_block = out_block[: -len(axis)] + np.unravel_index(out_block[-1], nblocks)
+            out_key = (self._name, *out_block)
+            in_key = (self.array.name, *in_block)
+            layer[out_key] = Alias(out_key, in_key)
+        return layer
+
+
+class FloxSubsetBlocks(ArrayExpr):
+    _parameters = ["array", "flatblocks", "blkshape", "reindexer", "output_chunks"]
+
+    @cached_property
+    def dtype(self):
+        return self.array.dtype
+
+    @cached_property
+    def chunks(self):
+        return self.output_chunks
+
+    @cached_property
+    def _meta(self):
+        return self.array._meta
+
+    def _layer(self):
+        index = _normalize_indexes(self.array.ndim, self.flatblocks, self.blkshape)
+        index = tuple(slice(k, k + 1) if isinstance(k, Integral) else k for k in index)
+
+        old_keys = np.empty(self.array.numblocks, dtype=object)
+        for block in itertools.product(*(range(n) for n in self.array.numblocks)):
+            old_keys[block] = (self.array.name, *block)
+
+        selected_keys = old_keys[index]
+        layer = {}
+        for out_block in itertools.product(*(range(len(chunks)) for chunks in self.output_chunks)):
+            in_key = selected_keys[out_block]
+            if isinstance(in_key, np.ndarray):
+                in_key = tuple(in_key.flat[0])
+            out_key = (self._name, *out_block)
+            layer[out_key] = Task(out_key, self.reindexer, TaskRef(in_key))
+        return layer
+
+
+def _unify_chunks(array, by, chunkmanager):
+    inds = tuple(range(array.ndim))
+
+    if not is_dask_array(by):
+        if is_duck_dask_array(by):
+            raise TypeError("Cannot mix dask_array.Array with other Dask-backed array types.")
+        chunks = tuple(array.chunks[ax] if by.shape[ax] != 1 else (1,) for ax in range(-by.ndim, 0))
+        by = chunkmanager.from_array(by, chunks=chunks)
+
+    _, (array, by) = chunkmanager.unify_chunks(array, inds, by, inds[-by.ndim :])
+    return array, by
+
+
+def _argreduce_preprocess(array, axis, chunkmanager):
+    assert len(axis) == 1
+    axis = axis[0]
+
+    idx = chunkmanager.array_api.arange(array.shape[axis], chunks=array.chunks[axis], dtype=np.intp)
+    idx = idx[tuple(slice(None) if i == axis else np.newaxis for i in range(array.ndim))]
+
+    def _zip_index(array_, idx_):
+        return (array_, idx_)
+
+    return chunkmanager.map_blocks(
+        _zip_index,
+        array,
+        idx,
+        dtype=array.dtype,
+        meta=array._meta,
+        name="groupby-argreduce-preprocess",
+    )
+
+
+def dask_groupby_agg(
+    array: Array,
+    by: Any,
+    *,
+    agg: Aggregation,
+    expected_groups: pd.RangeIndex | None,
+    reindex: ReindexStrategy,
+    axis: T_Axes = (),
+    fill_value: Any = None,
+    method: T_Method = "map-reduce",
+    engine: T_Engine = "numpy",
+    sort: bool = True,
+    chunks_cohorts=None,
+) -> tuple[Array, tuple[pd.Index | np.ndarray | Array]]:
+    from dask.array.core import slices_from_chunks
+
+    assert isinstance(axis, Sequence)
+    assert all(ax >= 0 for ax in axis)
+    if not is_dask_array(array):
+        raise TypeError("Cannot mix dask_array.Array with other Dask-backed array types.")
+
+    chunkmanager = get_chunkmanager(array, by)
+    array_api = chunkmanager.array_api
+
+    inds = tuple(range(array.ndim))
+    name = f"groupby_{agg.name}"
+
+    if expected_groups is None and reindex.blockwise:
+        raise ValueError("reindex.blockwise must be False-y if expected_groups is not provided.")
+    if method == "cohorts" and reindex.blockwise:
+        raise ValueError("reindex.blockwise must be False-y if method is 'cohorts'.")
+
+    by_input = by
+    array = new_collection(array.expr.simplify())
+    if is_dask_array(by):
+        by = new_collection(by.expr.simplify())
+    array, by = _unify_chunks(array, by, chunkmanager)
+    if method == "cohorts":
+        _, chunks_cohorts = find_group_cohorts(
+            by_input,
+            [array.chunks[ax] for ax in range(-by.ndim, 0)],
+            expected_groups=expected_groups,
+            merge=True,
+        )
+
+    token = dask.base.tokenize(array, by, agg, expected_groups, axis, method)
+
+    if agg.preprocess and method != "blockwise":
+        if _is_arg_reduction(agg):
+            array = _argreduce_preprocess(array, axis=axis, chunkmanager=chunkmanager)
+        else:
+            array = agg.preprocess(array, axis=axis)
+
+    labels_are_unknown = is_duck_dask_array(by_input) and expected_groups is None
+    do_grouped_combine = (
+        _is_arg_reduction(agg)
+        or labels_are_unknown
+        or (_is_first_last_reduction(agg) and array.dtype.kind != "f")
+    )
+    do_simple_combine = not do_grouped_combine
+
+    if method == "blockwise":
+        blockwise_method = partial(_reduce_blockwise, agg=agg, fill_value=fill_value, reindex=reindex)
+    else:
+        blockwise_method = partial(
+            _get_chunk_reduction(agg.reduction_type),
+            func=agg.chunk,
+            reindex=reindex.blockwise,
+            fill_value=agg.fill_value["intermediate"],
+            dtype=agg.dtype["intermediate"],
+            user_dtype=agg.dtype["user"],
+        )
+        if do_simple_combine:
+            blockwise_method = tlz.compose(_expand_dims, blockwise_method)
+
+    intermediate = chunkmanager.blockwise(
+        partial(
+            blockwise_method,
+            axis=axis,
+            expected_groups=expected_groups if reindex.blockwise else None,
+            engine=engine,
+            sort=sort,
+        ),
+        inds,
+        array,
+        inds,
+        by,
+        inds[-by.ndim :],
+        concatenate=False,
+        dtype=array.dtype,
+        meta=array._meta,
+        align_arrays=False,
+        name=f"{name}-chunk-{token}",
+    )
+
+    group_chunks: tuple[tuple[int | float, ...]]
+
+    if method in ["map-reduce", "cohorts"]:
+        combine = (
+            partial(_simple_combine, reindex=reindex)
+            if do_simple_combine
+            else partial(_grouped_combine, engine=engine, sort=sort)
+        )
+        aggregate = partial(_aggregate, combine=combine, agg=agg, fill_value=fill_value, reindex=reindex)
+
+        if method == "map-reduce":
+            reduced = array_api.reduction(
+                intermediate,
+                chunk=_passthrough_chunk,
+                aggregate=partial(aggregate, expected_groups=expected_groups),
+                axis=axis,
+                keepdims=True,
+                dtype=array.dtype,
+                combine=partial(combine, agg=agg),
+                name=f"{name}-simple-reduce",
+                concatenate=False,
+            )
+            if labels_are_unknown:
+                groups = (new_collection(FloxExtractGroups(reduced.expr, "groups", by.dtype)),)
+                group_chunks = ((np.nan,),)
+            else:
+                assert expected_groups is not None
+                groups = (expected_groups,)
+                group_chunks = ((len(expected_groups),),)
+
+        else:
+            assert chunks_cohorts
+            block_shape = intermediate.blocks.shape[-len(axis) :]
+            chunks_as_array = tuple(np.array(c) for c in intermediate.chunks)
+
+            cohort_results = []
+            groups_ = []
+            for icohort, (blks, cohort) in enumerate(chunks_cohorts.items()):
+                cohort_index = pd.Index(cohort)
+                reindexer = (
+                    partial(
+                        reindex_intermediates,
+                        agg=agg,
+                        unique_groups=cohort_index,
+                        array_type=reindex.array_type,
+                    )
+                    if do_simple_combine
+                    else identity
+                )
+
+                index = _normalize_indexes(intermediate.ndim, blks, block_shape)
+                index = tuple(slice(k, k + 1) if isinstance(k, Integral) else k for k in index)
+                squeezed = tuple(np.squeeze(i) if isinstance(i, np.ndarray) else i for i in index)
+                subset_chunks = tuple(tuple(c[i].tolist()) for c, i in zip(chunks_as_array, squeezed))
+
+                subset = new_collection(
+                    FloxSubsetBlocks(intermediate.expr, tuple(blks), block_shape, reindexer, subset_chunks)
+                )
+                new_reindex = ReindexStrategy(blockwise=do_simple_combine, array_type=reindex.array_type)
+                cohort_results.append(
+                    array_api.reduction(
+                        subset,
+                        chunk=_passthrough_chunk,
+                        aggregate=partial(
+                            aggregate,
+                            expected_groups=cohort_index,
+                            reindex=new_reindex,
+                            keepdims=True,
+                        ),
+                        axis=axis,
+                        keepdims=True,
+                        dtype=array.dtype,
+                        combine=partial(combine, agg=agg, reindex=new_reindex, keepdims=True),
+                        name=f"{name}-cohort-{icohort}-{token}",
+                        concatenate=False,
+                    )
+                )
+                groups_.append(cohort_index.values)
+
+            reduced = array_api.concatenate(cohort_results, axis=-1)
+            groups = (np.concatenate(groups_),)
+            group_chunks = (tuple(len(cohort) for cohort in groups_),)
+
+    elif method == "blockwise":
+        reduced = intermediate
+        if reindex.blockwise:
+            if TYPE_CHECKING:
+                assert expected_groups is not None
+            groups = (expected_groups,)
+            group_chunks = ((len(expected_groups),),)
+        else:
+            slices = slices_from_chunks(tuple(array.chunks[ax] for ax in axis))
+            from .core import _unique
+
+            groups_in_block = tuple(_unique(by_input[slc]) for slc in slices)
+            groups = (np.concatenate(groups_in_block),)
+            group_chunks = (tuple(len(grp) for grp in groups_in_block),)
+    else:
+        raise ValueError(f"Unknown method={method}.")
+
+    new_dims_shape = tuple(dim.size for dim in agg.new_dims if not dim.is_scalar)
+    new_inds = tuple(range(-len(new_dims_shape), 0))
+    out_inds = new_inds + inds[: -len(axis)] + (inds[-1],)
+    output_chunks = new_dims_shape + reduced.chunks[: -len(axis)] + group_chunks
+    new_axes = dict(zip(new_inds, new_dims_shape))
+
+    if method == "blockwise" and len(axis) > 1:
+        reduced = new_collection(FloxCollapseBlocks(reduced.expr, tuple(axis), group_chunks))
+
+    result = chunkmanager.blockwise(
+        _extract_result,
+        out_inds,
+        reduced,
+        inds,
+        adjust_chunks=dict(zip(out_inds, output_chunks)),
+        key=agg.name,
+        name=f"{name}-{token}",
+        concatenate=False,
+        new_axes=new_axes,
+        meta=reindex.get_dask_meta(array, dtype=agg.dtype["final"], fill_value=agg.fill_value[agg.name]),
+    )
+
+    return result, groups
+
+
+__all__ = [
+    "FloxCollapseBlocks",
+    "FloxExtractGroups",
+    "FloxSubsetBlocks",
+    "contains_dask_array",
+    "dask_groupby_agg",
+    "get_chunkmanager",
+    "is_dask_array",
+]
diff --git a/flox/factorize.py b/flox/factorize.py
index 4e66f8c7..439e8162 100644
--- a/flox/factorize.py
+++ b/flox/factorize.py
@@ -14,6 +14,7 @@
 import numpy as np
 import pandas as pd
 
+from .lib import contains_standalone_dask_array, is_standalone_dask_array
 from .types import FactorizeKwargs, FactorProps
 from .xrutils import is_duck_dask_array, isnull
 
@@ -234,10 +235,20 @@ def _factorize_multiple(
         sort=sort,
     )
     if any_by_dask:
-        import dask.array
-
         from . import dask_array_ops  # noqa
 
+        dask_bys = tuple(by_ for by_ in by if is_duck_dask_array(by_))
+        if contains_standalone_dask_array(*dask_bys):
+            if not all(is_standalone_dask_array(by_) for by_ in dask_bys):
+                raise TypeError("Cannot mix dask_array.Array with other Dask-backed array types.")
+            from .dask_array import get_chunkmanager
+
+            chunkmanager = get_chunkmanager(*by)
+        else:
+            import dask.array
+
+            chunkmanager = dask.array
+
         # unifying chunks will make sure all arrays in `by` are dask arrays
         # with compatible chunks, even if there was originally a numpy array
         inds = tuple(range(by[0].ndim))
@@ -251,9 +262,9 @@ def _factorize_multiple(
         )
         grp_shape = tuple(map(len, found_groups))
 
-        chunks, by_chunked = dask.array.unify_chunks(*itertools.chain(*zip(by, (inds,) * len(by))))
+        chunks, by_chunked = chunkmanager.unify_chunks(*itertools.chain(*zip(by, (inds,) * len(by))))
         group_idxs = [
-            dask.array.map_blocks(
+            chunkmanager.map_blocks(
                 _lazy_factorize_wrapper,
                 by_,
                 expected_groups=(expect_,),
@@ -264,7 +275,7 @@ def _factorize_multiple(
         ]
         # This could be avoied but we'd use `np.where`
         # instead `_ravel_factorized` instead i.e. a copy.
-        group_idx = dask.array.map_blocks(
+        group_idx = chunkmanager.map_blocks(
             _ravel_factorized, *group_idxs, grp_shape=grp_shape, chunks=tuple(chunks.values()), dtype=np.int64
         )
 
diff --git a/flox/lib.py b/flox/lib.py
index 2c2c1a51..2bf9e926 100644
--- a/flox/lib.py
+++ b/flox/lib.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, TypeAlias, TypeVar
+from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar
 
 from .types import DaskArray, Graph
 from .xrutils import is_duck_dask_array, module_available
@@ -13,12 +13,13 @@
 
 T = TypeVar("T")
 
+dask_array_type: tuple[type[Any], ...] = ()
 try:
     import dask.array as da
 
-    dask_array_type = da.Array
+    dask_array_type += (da.Array,)
 except ImportError:
-    dask_array_type = ()  # type: ignore[assignment, misc]
+    pass
 
 try:
     import sparse
@@ -30,6 +31,14 @@
 HAS_SPARSE = module_available("sparse")
 
 
+def is_standalone_dask_array(x) -> bool:
+    return type(x).__module__.split(".", 1)[0] == "dask_array" and is_duck_dask_array(x)
+
+
+def contains_standalone_dask_array(*args) -> bool:
+    return any(is_standalone_dask_array(arg) for arg in args)
+
+
 @dataclass
 class ArrayLayer:
     name: str
diff --git a/flox/rechunk.py b/flox/rechunk.py
index 578991fd..018ab394 100644
--- a/flox/rechunk.py
+++ b/flox/rechunk.py
@@ -16,6 +16,7 @@
 from .aggregations import _atleast_1d
 from .cache import memoize
 from .factorize import factorize_
+from .lib import is_standalone_dask_array
 from .options import OPTIONS
 
 if TYPE_CHECKING:
@@ -181,6 +182,11 @@ def rechunk_for_blockwise(
         Rechunked array
     """
 
+    if is_standalone_dask_array(array):
+        from dask_array._new_collection import new_collection
+
+        array = new_collection(array.expr.simplify())
+
     chunks = array.chunks[axis]
     if len(chunks) == 1:
         return "blockwise", array
diff --git a/flox/reindex.py b/flox/reindex.py
index b70c444a..f1c6cd84 100644
--- a/flox/reindex.py
+++ b/flox/reindex.py
@@ -13,8 +13,8 @@
 import pandas as pd
 
 from . import xrdtypes
-from .lib import dask_array_type, sparse_array_type
-from .xrutils import isnull
+from .lib import sparse_array_type
+from .xrutils import is_duck_dask_array, isnull
 
 if TYPE_CHECKING:
     from .core import T_Axis
@@ -77,7 +77,7 @@ def set_blockwise_for_numpy(self):
 
     def get_dask_meta(self, other, *, fill_value, dtype) -> Any:
         if self.array_type is ReindexArrayType.AUTO:
-            other = other._meta if isinstance(other, dask_array_type) else other
+            other = other._meta if is_duck_dask_array(other) else other
             if isinstance(other, sparse_array_type):
                 return type(other).from_numpy(np.array([], dtype=dtype))
             return type(other)([], dtype=dtype)
diff --git a/flox/xrutils.py b/flox/xrutils.py
index 7edc63f6..685c2385 100644
--- a/flox/xrutils.py
+++ b/flox/xrutils.py
@@ -53,10 +53,11 @@ def module_available(module: str, minversion: str | None = None) -> bool:
     cubed = None
 
 dask: ModuleType | None
+dask_array_type: tuple[type[Any], ...]
 try:
     import dask.array
 
-    dask_array_type = dask.array.Array  # type: ignore[union-attr]
+    dask_array_type = (dask.array.Array,)  # type: ignore[union-attr]
 except ImportError:
     dask = None
     dask_array_type = ()
diff --git a/pyproject.toml b/pyproject.toml
index 59f73026..822d4a7d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -53,6 +53,7 @@ docs = [
     "sphinx-codeautolink",
     "sphinx-copybutton",
     "pyarrow",
+    "rich",
 ]
 
 [build-system]
@@ -146,6 +147,8 @@ module=[
     "cftime",
     "cubed.*",
     "dask.*",
+    "dask_array",
+    "dask_array.*",
     "importlib_metadata",
     "numba",
     "numbagg.*",
@@ -254,6 +257,7 @@ docs = [
     "jupyter",
     "sphinx-codeautolink",
     "sphinx-copybutton",
+    "rich",
 ]
 benchmark = [
     "asv>=0.6.4",
diff --git a/tests/__init__.py b/tests/__init__.py
index a27f3855..57d61b08 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -1,4 +1,5 @@
 import importlib
+import importlib.util
 from contextlib import nullcontext
 from typing import Any
 
@@ -7,7 +8,7 @@
 import pandas as pd
 import pytest
 
-from flox.lib import dask_array_type, sparse_array_type
+from flox.lib import sparse_array_type
 from flox.xrutils import is_duck_dask_array
 
 pd_types = (pd.Index,)
@@ -33,6 +34,12 @@ def _importorskip(modname, minversion=None):
     return has, func
 
 
+def _module_available_without_import(modname):
+    has = importlib.util.find_spec(modname) is not None
+    func = pytest.mark.skipif(not has, reason=f"requires {modname}")
+    return has, func
+
+
 def LooseVersion(vstring):
     # Our development version is something like '0.10.9+aac7bfc'
     # This function just ignored the git commit id.
@@ -43,6 +50,7 @@ def LooseVersion(vstring):
 has_cftime, requires_cftime = _importorskip("cftime")
 has_cubed, requires_cubed = _importorskip("cubed")
 has_dask, requires_dask = _importorskip("dask")
+has_dask_array, requires_dask_array = _module_available_without_import("dask_array")
 has_sparse, requires_sparse = _importorskip("sparse")
 has_numba, requires_numba = _importorskip("numba")
 has_numbagg, requires_numbagg = _importorskip("numbagg")
@@ -105,11 +113,10 @@ def assert_equal(a, b, tolerance=None):
     if a.dtype != b.dtype:
         raise AssertionError(f"a and b have different dtypes: (a: {a.dtype}, b: {b.dtype})")
 
-    if has_dask:
-        a_eager = a.compute() if isinstance(a, dask_array_type) else a
-        b_eager = b.compute() if isinstance(b, dask_array_type) else b
-    else:
-        a_eager, b_eager = a, b
+    a_is_chunked = is_duck_dask_array(a)
+    b_is_chunked = is_duck_dask_array(b)
+    a_eager = a.compute() if a_is_chunked else a
+    b_eager = b.compute() if b_is_chunked else b
 
     if has_sparse:
         one_is_sparse = isinstance(a_eager, sparse_array_type) or isinstance(b_eager, sparse_array_type)
@@ -123,7 +130,7 @@ def assert_equal(a, b, tolerance=None):
     else:
         np.testing.assert_allclose(a_eager, b_eager, equal_nan=True, **tolerance)
 
-    if has_dask and isinstance(a, dask_array_type) or isinstance(b, dask_array_type):
+    if a_is_chunked or b_is_chunked:
         # does some validation of the dask graph
         dask_assert_eq(a, b, equal_nan=True, check_type=not one_is_sparse)
 
diff --git a/tests/conftest.py b/tests/conftest.py
index ac0624ce..5954f337 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,7 +1,7 @@
 import pytest
 from hypothesis import HealthCheck, Verbosity, settings
 
-from . import requires_numbagg
+from . import has_dask_array, requires_numbagg
 
 settings.register_profile(
     "ci",
@@ -30,3 +30,73 @@
 )
 def engine(request):
     return request.param
+
+
+_MISSING = object()
+
+
+def _snapshot_xarray_dask_manager():
+    try:
+        from xarray.namedarray.parallelcompat import list_chunkmanagers
+    except ImportError:
+        return None, _MISSING
+
+    managers = list_chunkmanagers()
+    return managers, managers.get("dask", _MISSING)
+
+
+def _restore_xarray_dask_manager(managers, original):
+    if managers is None:
+        return
+    if original is _MISSING:
+        managers.pop("dask", None)
+    else:
+        managers["dask"] = original
+
+
+def _activate_dask_array():
+    import dask_array as da
+
+    da.xarray.register()
+    return da
+
+
+@pytest.fixture
+def dask_array_api():
+    if not has_dask_array:
+        pytest.skip("requires dask_array")
+
+    managers, original = _snapshot_xarray_dask_manager()
+    da = _activate_dask_array()
+
+    try:
+        yield da
+    finally:
+        _restore_xarray_dask_manager(managers, original)
+
+
+@pytest.fixture
+def chunked_array_api(request):
+    try:
+        backend = request.param
+    except AttributeError:
+        raise ValueError("chunked_array_api must be parametrized indirectly") from None
+
+    if backend == "dask":
+        import dask.array as da
+
+        yield da
+        return
+
+    if backend != "dask_array":
+        raise ValueError(f"Unknown chunked array backend {backend!r}")
+
+    if not has_dask_array:
+        pytest.skip("requires dask_array")
+
+    managers, original = _snapshot_xarray_dask_manager()
+    da = _activate_dask_array()
+    try:
+        yield da
+    finally:
+        _restore_xarray_dask_manager(managers, original)
diff --git a/tests/test_core.py b/tests/test_core.py
index 2154b991..e877c90d 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -45,6 +45,7 @@
     raise_if_dask_computes,
     requires_cubed,
     requires_dask,
+    requires_dask_array,
     requires_sparse,
     to_numpy,
 )
@@ -78,6 +79,10 @@ def dask_array_ones(*args):
 REINDEX_SPARSE_PARAM = pytest.param(
     REINDEX_SPARSE_STRAT, marks=(requires_dask, pytest.mark.skipif(not has_sparse, reason="no sparse"))
 )
+CHUNKED_ARRAY_BACKENDS = (
+    pytest.param("dask", marks=requires_dask, id="dask"),
+    pytest.param("dask_array", marks=requires_dask_array, id="dask-array"),
+)
 
 if TYPE_CHECKING:
     from flox.core import T_Agg, T_Engine, T_ExpectedGroupsOpt, T_Method
@@ -468,7 +473,7 @@ def test_numpy_reduce_nd_md():
     assert_equal(actual, expected)
 
 
-@requires_dask
+@pytest.mark.parametrize("chunked_array_api", CHUNKED_ARRAY_BACKENDS, indirect=True)
 @pytest.mark.parametrize("reindex", [None, False, True, REINDEX_SPARSE_PARAM])
 @pytest.mark.parametrize("func", ALL_FUNCS)
 @pytest.mark.parametrize("add_nan", [False, True])
@@ -483,7 +488,9 @@ def test_numpy_reduce_nd_md():
         ((10, 12), (3, 3), 3),  # form 3
     ],
 )
-def test_groupby_agg_dask(func, shape, array_chunks, group_chunks, add_nan, dtype, engine, reindex):
+def test_groupby_agg_dask(
+    func, shape, array_chunks, group_chunks, add_nan, dtype, engine, reindex, chunked_array_api
+):
     """Tests groupby_reduce with dask arrays against groupby_reduce with numpy arrays"""
 
     if func in ["first", "last"] or func in BLOCKWISE_FUNCS:
@@ -496,8 +503,8 @@ def test_groupby_agg_dask(func, shape, array_chunks, group_chunks, add_nan, dtyp
         pytest.skip()
 
     rng = np.random.default_rng(12345)
-    array = dask.array.from_array(rng.random(shape), chunks=array_chunks).astype(dtype)
-    array = dask.array.ones(shape, chunks=array_chunks)
+    array = chunked_array_api.from_array(rng.random(shape), chunks=array_chunks).astype(dtype)
+    array = chunked_array_api.ones(shape, chunks=array_chunks)
 
     labels = np.array([0, 0, 2, 2, 2, 1, 1, 2, 2, 1, 1, 0])
     if add_nan:
@@ -519,7 +526,7 @@ def test_groupby_agg_dask(func, shape, array_chunks, group_chunks, add_nan, dtyp
         actual, _ = groupby_reduce(array, labels, engine=engine, **kwargs)
     assert_equal(actual, expected)
 
-    by = from_array(labels, group_chunks)
+    by = chunked_array_api.from_array(labels, group_chunks)
     with raise_if_dask_computes():
         actual, _ = groupby_reduce(array, by, engine=engine, **kwargs)
     assert_equal(expected, actual)
@@ -611,14 +618,14 @@ def test_numpy_reduce_axis_subset(engine):
     assert_equal(result, expected)
 
 
-@requires_dask
-def test_dask_reduce_axis_subset():
+@pytest.mark.parametrize("chunked_array_api", CHUNKED_ARRAY_BACKENDS, indirect=True)
+def test_dask_reduce_axis_subset(chunked_array_api):
     by = labels2d
     array = np.ones_like(by, dtype=np.int64)
     with raise_if_dask_computes():
         result, _ = groupby_reduce(
-            da.from_array(array, chunks=(2, 3)),
-            da.from_array(by, chunks=(2, 2)),
+            chunked_array_api.from_array(array, chunks=(2, 3)),
+            chunked_array_api.from_array(by, chunks=(2, 2)),
             func="count",
             axis=1,
             expected_groups=[0, 2],
@@ -631,8 +638,8 @@ def test_dask_reduce_axis_subset():
     expected = np.tile(subarr, (3, 1, 1))
     with raise_if_dask_computes():
         result, _ = groupby_reduce(
-            da.from_array(array, chunks=(1, 2, 3)),
-            da.from_array(by, chunks=(2, 2, 2)),
+            chunked_array_api.from_array(array, chunks=(1, 2, 3)),
+            chunked_array_api.from_array(by, chunks=(2, 2, 2)),
             func="count",
             axis=1,
             expected_groups=[0, 2],
@@ -644,8 +651,8 @@ def test_dask_reduce_axis_subset():
     expected = np.tile(subarr, (3, 1, 1))
     with raise_if_dask_computes():
         result, _ = groupby_reduce(
-            da.from_array(array, chunks=(1, 2, 3)),
-            da.from_array(by, chunks=(2, 2, 2)),
+            chunked_array_api.from_array(array, chunks=(1, 2, 3)),
+            chunked_array_api.from_array(by, chunks=(2, 2, 2)),
             func="count",
             axis=2,
             expected_groups=[0, 2],
@@ -654,8 +661,8 @@ def test_dask_reduce_axis_subset():
 
     with pytest.raises(NotImplementedError):
         groupby_reduce(
-            da.from_array(array, chunks=(1, 3, 2)),
-            da.from_array(by, chunks=(2, 2, 2)),
+            chunked_array_api.from_array(array, chunks=(1, 3, 2)),
+            chunked_array_api.from_array(by, chunks=(2, 2, 2)),
             func="count",
             axis=2,
         )
@@ -938,27 +945,34 @@ def test_npg_nanarg_bug(func):
     ),
 )
 @pytest.mark.parametrize("method", ["cohorts", "map-reduce"])
-@pytest.mark.parametrize("chunk_labels", [False, True])
+def test_groupby_bins(kwargs, engine, method) -> None:
+    _assert_groupby_bins([1, 1, 1, 1, 1, 1], [0.2, 1.5, 1.9, 2, 3, 20], kwargs, engine, method)
+
+
 @pytest.mark.parametrize(
-    "chunks",
+    "kwargs",
     (
-        (),
-        pytest.param((1,), marks=requires_dask),
-        pytest.param((2,), marks=requires_dask),
+        dict(expected_groups=np.array([1, 2, 4, 5]), isbin=True),
+        dict(expected_groups=pd.IntervalIndex.from_breaks([1, 2, 4, 5])),
     ),
 )
-def test_groupby_bins(chunk_labels, kwargs, chunks, engine, method) -> None:
-    array = [1, 1, 1, 1, 1, 1]
-    labels = [0.2, 1.5, 1.9, 2, 3, 20]
-
+@pytest.mark.parametrize("method", ["cohorts", "map-reduce"])
+@pytest.mark.parametrize("chunk_labels", [False, True])
+@pytest.mark.parametrize("chunks", [(1,), (2,)])
+@pytest.mark.parametrize("chunked_array_api", CHUNKED_ARRAY_BACKENDS, indirect=True)
+def test_groupby_bins_chunked(chunked_array_api, chunk_labels, kwargs, chunks, engine, method) -> None:
     if method == "cohorts" and chunk_labels:
         pytest.xfail()
 
-    if chunks:
-        array = dask.array.from_array(array, chunks=chunks)
-        if chunk_labels:
-            labels = dask.array.from_array(labels, chunks=chunks)
+    array = chunked_array_api.from_array([1, 1, 1, 1, 1, 1], chunks=chunks)
+    labels = [0.2, 1.5, 1.9, 2, 3, 20]
+    if chunk_labels:
+        labels = chunked_array_api.from_array(labels, chunks=chunks)
+
+    _assert_groupby_bins(array, labels, kwargs, engine, method)
 
+
+def _assert_groupby_bins(array, labels, kwargs, engine, method) -> None:
     with raise_if_dask_computes():
         actual, *groups = groupby_reduce(
             array,
@@ -1132,7 +1146,6 @@ def test_fill_value_behaviour(func, chunks, fill_value, engine):
     assert_equal(actual, expected)
 
 
-@requires_dask
 @pytest.mark.parametrize("func", ["mean", "sum"])
 @pytest.mark.parametrize("dtype", ["float32", "float64", "int32", "int64"])
 def test_dtype_preservation(dtype, func, engine):
@@ -1150,19 +1163,35 @@ def test_dtype_preservation(dtype, func, engine):
     actual, _ = groupby_reduce(array, by, func=func, engine=engine)
     assert actual.dtype == expected
 
-    array = dask.array.from_array(array, chunks=(4,))
+
+@pytest.mark.parametrize("chunked_array_api", CHUNKED_ARRAY_BACKENDS, indirect=True)
+@pytest.mark.parametrize("func", ["mean", "sum"])
+@pytest.mark.parametrize("dtype", ["float32", "float64", "int32", "int64"])
+def test_dtype_preservation_chunked(chunked_array_api, dtype, func, engine):
+    if engine == "numbagg":
+        # https://github.com/numbagg/numbagg/issues/121
+        pytest.skip()
+    if func == "sum":
+        expected = dtypes._maybe_promote_int(dtype)
+    elif func == "mean" and "int" in dtype:
+        expected = np.float64
+    else:
+        expected = np.dtype(dtype)
+    array = np.ones((20,), dtype=dtype)
+    by = np.ones(array.shape, dtype=int)
+    array = chunked_array_api.from_array(array, chunks=(4,))
     actual, _ = groupby_reduce(array, by, func=func, engine=engine)
     assert actual.dtype == expected
 
 
-@requires_dask
+@pytest.mark.parametrize("chunked_array_api", CHUNKED_ARRAY_BACKENDS, indirect=True)
 @pytest.mark.parametrize("dtype", [np.float32, np.float64, np.int32, np.int64])
 @pytest.mark.parametrize("labels_dtype", [np.float32, np.float64, np.int32, np.int64])
 @pytest.mark.parametrize("method", ["map-reduce", "cohorts"])
-def test_cohorts_map_reduce_consistent_dtypes(method, dtype, labels_dtype):
+def test_cohorts_map_reduce_consistent_dtypes(method, dtype, labels_dtype, chunked_array_api):
     repeats = np.array([4, 4, 12, 2, 3, 4], dtype=np.int32)
     labels = np.repeat(np.arange(6, dtype=labels_dtype), repeats)
-    array = dask.array.from_array(labels.astype(dtype), chunks=(4, 8, 4, 9, 4))
+    array = chunked_array_api.from_array(labels.astype(dtype), chunks=(4, 8, 4, 9, 4))
 
     actual, actual_groups = groupby_reduce(array, labels, func="count", method=method)
     assert_equal(actual_groups, np.arange(6, dtype=labels.dtype))
@@ -1175,11 +1204,11 @@ def test_cohorts_map_reduce_consistent_dtypes(method, dtype, labels_dtype):
     assert_equal(actual, np.array([0, 4, 24, 6, 12, 20], dtype=expect_dtype))
 
 
-@requires_dask
+@pytest.mark.parametrize("chunked_array_api", CHUNKED_ARRAY_BACKENDS, indirect=True)
 @pytest.mark.parametrize("func", ALL_FUNCS)
 @pytest.mark.parametrize("axis", (-1, None))
 @pytest.mark.parametrize("method", ["blockwise", "cohorts", "map-reduce"])
-def test_cohorts_nd_by(func, method, axis, engine):
+def test_cohorts_nd_by(func, method, axis, engine, chunked_array_api):
     if (
         ("arg" in func and (axis is None or engine in ["flox", "numbagg"]))
         or (method != "blockwise" and func in BLOCKWISE_FUNCS)
@@ -1189,10 +1218,10 @@ def test_cohorts_nd_by(func, method, axis, engine):
     if axis is not None and method != "map-reduce":
         pytest.xfail()
 
-    o = dask.array.ones((3,), chunks=-1)
-    o2 = dask.array.ones((2, 3), chunks=-1)
+    o = chunked_array_api.ones((3,), chunks=-1)
+    o2 = chunked_array_api.ones((2, 3), chunks=-1)
 
-    array = dask.array.block([[o, 2 * o], [3 * o2, 4 * o2]])
+    array = chunked_array_api.block([[o, 2 * o], [3 * o2, 4 * o2]])
     by = array.compute().astype(np.int64)
     by[0, 1] = 30
     by[2, 1] = 40
@@ -1381,16 +1410,22 @@ def test_factorize_values_outside_bins():
     assert_equal(expected, actual)
 
 
-@pytest.mark.parametrize("chunk", [pytest.param(True, marks=requires_dask), False])
-def test_multiple_groupers_bins(chunk) -> None:
-    xp = dask.array if chunk else np
-    array_kwargs = {"chunks": 2} if chunk else {}
-    array = xp.ones((5, 2), **array_kwargs, dtype=np.int64)
+def test_multiple_groupers_bins() -> None:
+    _assert_multiple_groupers_bins(np, {})
+
+
+@pytest.mark.parametrize("chunked_array_api", CHUNKED_ARRAY_BACKENDS, indirect=True)
+def test_multiple_groupers_bins_chunked(chunked_array_api) -> None:
+    _assert_multiple_groupers_bins(chunked_array_api, {"chunks": 2})
+
+
+def _assert_multiple_groupers_bins(array_api, array_kwargs) -> None:
+    array = array_api.ones((5, 2), **array_kwargs, dtype=np.int64)
 
     actual, *_ = groupby_reduce(
         array,
         np.arange(10).reshape(5, 2),
-        xp.arange(10).reshape(5, 2),
+        array_api.arange(10).reshape(5, 2),
         axis=(0, 1),
         expected_groups=(
             pd.IntervalIndex.from_breaks(np.arange(2, 8, 1)),
@@ -1413,17 +1448,28 @@ def test_multiple_groupers_bins(chunk) -> None:
         np.arange(2, 4).reshape(1, 2),
     ],
 )
-@pytest.mark.parametrize("chunk", [pytest.param(True, marks=requires_dask), False])
-def test_multiple_groupers(chunk, by1, by2, expected_groups) -> None:
-    if chunk and expected_groups is None:
-        pytest.skip()
+def test_multiple_groupers(by1, by2, expected_groups) -> None:
+    _assert_multiple_groupers(np, {}, by1, by2, expected_groups)
 
-    xp = dask.array if chunk else np
-    array_kwargs = {"chunks": 2} if chunk else {}
-    array = xp.ones((5, 2), **array_kwargs, dtype=np.int64)
 
-    if chunk:
-        by2 = dask.array.from_array(by2)
+@pytest.mark.parametrize("expected_groups", [(np.arange(5), [2, 3]), (None, [2, 3])])
+@pytest.mark.parametrize("by1", [np.arange(5)[:, None], np.broadcast_to(np.arange(5)[:, None], (5, 2))])
+@pytest.mark.parametrize(
+    "by2",
+    [
+        np.arange(2, 4).reshape(1, 2),
+        np.broadcast_to(np.arange(2, 4).reshape(1, 2), (5, 2)),
+        np.arange(2, 4).reshape(1, 2),
+    ],
+)
+@pytest.mark.parametrize("chunked_array_api", CHUNKED_ARRAY_BACKENDS, indirect=True)
+def test_multiple_groupers_chunked(chunked_array_api, by1, by2, expected_groups) -> None:
+    by2 = chunked_array_api.from_array(by2, chunks=by2.shape)
+    _assert_multiple_groupers(chunked_array_api, {"chunks": 2}, by1, by2, expected_groups)
+
+
+def _assert_multiple_groupers(array_api, array_kwargs, by1, by2, expected_groups) -> None:
+    array = array_api.ones((5, 2), **array_kwargs, dtype=np.int64)
 
     # output from `count` is intp
     expected = np.ones((5, 2), dtype=np.intp)
@@ -1449,13 +1495,13 @@ def test_validate_expected_groups(expected_groups):
         )
 
 
-@requires_dask
-def test_validate_expected_groups_not_none_dask() -> None:
+@pytest.mark.parametrize("chunked_array_api", CHUNKED_ARRAY_BACKENDS, indirect=True)
+def test_validate_expected_groups_not_none_dask(chunked_array_api) -> None:
     with pytest.raises(ValueError):
         groupby_reduce(
-            dask.array.ones((5, 2)),
+            chunked_array_api.ones((5, 2)),
             np.arange(10).reshape(5, 2),
-            dask.array.arange(10).reshape(5, 2),
+            chunked_array_api.arange(10).reshape(5, 2),
             axis=(0, 1),
             expected_groups=None,
             func="count",
diff --git a/tests/test_dask_array_expr.py b/tests/test_dask_array_expr.py
new file mode 100644
index 00000000..d949c611
--- /dev/null
+++ b/tests/test_dask_array_expr.py
@@ -0,0 +1,211 @@
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+
+import numpy as np
+import pytest
+
+from flox.core import groupby_reduce
+
+from . import requires_dask
+
+
+def _contains_expr(expr, name, seen=None):
+    if seen is None:
+        seen = set()
+    if id(expr) in seen:
+        return False
+    seen.add(id(expr))
+    if type(expr).__name__ == name:
+        return True
+    if not hasattr(expr, "dependencies"):
+        return False
+    return any(_contains_expr(getattr(dep, "expr", dep), name, seen) for dep in expr.dependencies())
+
+
+def test_groupby_reduce_with_dask_array_returns_dask_array(dask_array_api):
+    dax = dask_array_api
+    x = dax.from_array(np.arange(6), chunks=(3,))
+    labels = np.array([0, 0, 1, 1, 2, 2])
+
+    result, groups = groupby_reduce(
+        x,
+        labels,
+        func="sum",
+        expected_groups=np.array([0, 1, 2]),
+    )
+
+    assert isinstance(result, dax.Array)
+    np.testing.assert_array_equal(groups, np.array([0, 1, 2]))
+    np.testing.assert_array_equal(result.compute(), np.array([1, 5, 9]))
+    assert "FromGraph" not in type(result.expr).__name__
+
+
+def test_groupby_reduce_with_dask_array_labels(dask_array_api):
+    dax = dask_array_api
+    x = dax.from_array(np.arange(6), chunks=(3,))
+    labels = dax.from_array(np.array([0, 0, 1, 1, 2, 2]), chunks=(3,))
+
+    result, groups = groupby_reduce(
+        x,
+        labels,
+        func="sum",
+        expected_groups=np.array([0, 1, 2]),
+    )
+
+    assert isinstance(result, dax.Array)
+    np.testing.assert_array_equal(groups, np.array([0, 1, 2]))
+    np.testing.assert_array_equal(result.compute(), np.array([1, 5, 9]))
+
+
+def test_groupby_reduce_with_dask_array_unknown_groups_uses_expression(dask_array_api):
+    dax = dask_array_api
+    x = dax.from_array(np.arange(6), chunks=(3,))
+    labels = dax.from_array(np.array([0, 0, 1, 1, 2, 2]), chunks=(3,))
+
+    result, groups = groupby_reduce(x, labels, func="sum")
+
+    assert isinstance(result, dax.Array)
+    assert isinstance(groups, dax.Array)
+    assert type(groups.expr).__name__ == "FloxExtractGroups"
+    np.testing.assert_array_equal(groups.compute(), np.array([0, 1, 2]))
+    np.testing.assert_array_equal(result.compute(), np.array([1, 5, 9]))
+
+
+def test_groupby_reduce_with_dask_array_cohorts_uses_subset_expression(dask_array_api):
+    dax = dask_array_api
+    x = dax.from_array(np.arange(12).reshape(3, 4), chunks=(1, 2))
+    labels = np.array([[0, 0, 1, 1], [0, 0, 1, 1], [2, 2, 3, 3]])
+
+    result, groups = groupby_reduce(
+        x,
+        labels,
+        func="sum",
+        method="cohorts",
+        expected_groups=np.array([0, 1, 2, 3]),
+        axis=(0, 1),
+    )
+
+    assert isinstance(result, dax.Array)
+    assert _contains_expr(result.expr, "FloxSubsetBlocks")
+    assert not _contains_expr(result.expr, "FromGraph")
+    np.testing.assert_array_equal(groups, np.array([0, 1, 2, 3]))
+    np.testing.assert_array_equal(result.compute(), np.array([10, 18, 17, 21]))
+
+
+def test_xarray_reduce_with_dask_array_data(dask_array_api):
+    dax = dask_array_api
+    xr = pytest.importorskip("xarray")
+    from flox.xarray import xarray_reduce
+
+    x = xr.DataArray(dax.from_array(np.arange(6), chunks=(3,)), dims="x")
+    labels = xr.DataArray(np.array([0, 0, 1, 1, 2, 2]), dims="x", name="group")
+
+    result = xarray_reduce(x, labels, func="sum")
+
+    assert isinstance(result.data, dax.Array)
+    np.testing.assert_array_equal(result.compute().values, np.array([1, 5, 9]))
+
+
+@pytest.mark.parametrize("method", ["xarray", "map-reduce", "cohorts", "blockwise"])
+def test_xarray_reduce_with_dask_array_rolling_groupby(method, dask_array_api):
+    dax = dask_array_api
+    xr = pytest.importorskip("xarray")
+    from flox.xarray import xarray_reduce
+
+    time = np.arange(
+        np.datetime64("2015-12-30T00"),
+        np.datetime64("2016-01-03T00"),
+        np.timedelta64(1, "h"),
+    )
+    data = np.arange(time.size * 2 * 3, dtype=np.float32).reshape(time.size, 2, 3)
+    x = xr.DataArray(
+        dax.from_array(data, chunks=(6, 2, 3)),
+        dims=("time", "lat", "lon"),
+        coords={"time": time},
+    )
+    expected = (
+        xr.DataArray(data, dims=("time", "lat", "lon"), coords={"time": time})
+        .rolling(time=3, min_periods=3)
+        .sum()
+        .groupby("time.year")
+        .max("time")
+    )
+
+    rolled = x.rolling(time=3, min_periods=3).sum()
+    labels = xr.DataArray(
+        rolled.time.dt.year.data,
+        name="year",
+        dims="time",
+        coords={"time": rolled.time},
+    )
+
+    if method == "xarray":
+        result = rolled.groupby("time.year").max("time")
+    else:
+        result = xarray_reduce(
+            rolled,
+            labels,
+            func="max",
+            dim="time",
+            expected_groups=np.array([2015, 2016]),
+            method=method,
+        )
+
+    assert isinstance(result.data, dax.Array)
+    expected_chunks = ((2,), (2,), (3,)) if method in ("xarray", "map-reduce") else ((1, 1), (2,), (3,))
+    assert result.data.chunks == expected_chunks
+    xr.testing.assert_identical(result.compute().rename(None), expected)
+
+
+@requires_dask
+def test_importing_flox_does_not_register_dask_array_for_legacy_xarray_dask():
+    code = """
+import sys
+import numpy as np
+import dask.array as da
+import xarray as xr
+from flox.xarray import xarray_reduce
+
+assert "dask_array" not in sys.modules
+x = xr.DataArray(da.from_array(np.arange(6), chunks=(3,)), dims="x")
+labels = xr.DataArray(np.array([0, 0, 1, 1, 2, 2]), dims="x", name="group")
+result = xarray_reduce(x, labels, func="sum")
+assert "dask_array" not in sys.modules
+np.testing.assert_array_equal(result.compute().values, np.array([1, 5, 9]))
+"""
+    subprocess.run(
+        [sys.executable, "-c", code],
+        env=os.environ.copy(),
+        check=True,
+        text=True,
+        capture_output=True,
+    )
+
+
+def test_mixed_dask_backends_raise_clear_error(dask_array_api):
+    dax = dask_array_api
+    legacy_da = pytest.importorskip("dask.array")
+    x = dax.from_array(np.arange(6), chunks=(3,))
+    labels = legacy_da.from_array(np.array([0, 0, 1, 1, 2, 2]), chunks=(3,))
+
+    with pytest.raises(TypeError, match="Cannot mix dask_array.Array"):
+        groupby_reduce(
+            x,
+            labels,
+            func="sum",
+            expected_groups=np.array([0, 1, 2]),
+        )
+
+    legacy_x = legacy_da.from_array(np.arange(6), chunks=(3,))
+    dax_labels = dax.from_array(np.array([0, 0, 1, 1, 2, 2]), chunks=(3,))
+
+    with pytest.raises(TypeError, match="Cannot mix dask_array.Array"):
+        groupby_reduce(
+            legacy_x,
+            dax_labels,
+            func="sum",
+            expected_groups=np.array([0, 1, 2]),
+        )
diff --git a/tests/test_xarray.py b/tests/test_xarray.py
index 798ef847..b7126e2f 100644
--- a/tests/test_xarray.py
+++ b/tests/test_xarray.py
@@ -17,6 +17,7 @@
     raise_if_dask_computes,
     requires_cftime,
     requires_dask,
+    requires_dask_array,
 )
 
 if has_dask:
@@ -29,6 +30,10 @@
 xr.set_options(use_flox=False, use_numbagg=False, use_bottleneck=False)
 tolerance64 = {"rtol": 1e-15, "atol": 1e-18}
 np.random.seed(123)
+CHUNKED_ARRAY_BACKENDS = (
+    pytest.param("dask", marks=requires_dask, id="dask"),
+    pytest.param("dask_array", marks=requires_dask_array, id="dask-array"),
+)
 
 
 @pytest.mark.parametrize("reindex", [None, False, True])
@@ -214,14 +219,14 @@ def test_xarray_reduce_cftime_var(engine, indexer, expected_groups, func):
 
 
 @requires_cftime
-@requires_dask
-def test_xarray_reduce_single_grouper(engine):
+@pytest.mark.parametrize("chunked_array_api", CHUNKED_ARRAY_BACKENDS, indirect=True)
+def test_xarray_reduce_single_grouper(engine, chunked_array_api):
     # DataArray
     ds = xr.Dataset(
         {
             "Tair": (
                 ("time", "x", "y"),
-                dask.array.ones((36, 205, 275), chunks=(9, -1, -1)),
+                chunked_array_api.ones((36, 205, 275), chunks=(9, -1, -1)),
             )
         },
         coords={"time": xr.date_range("1980-09-01 00:00", "1983-09-18 00:00", freq="ME", calendar="noleap")},
@@ -395,15 +400,15 @@ def test_xarray_groupby_bins(chunks, engine):
     xr.testing.assert_equal(actual, expected)
 
 
-@requires_dask
-def test_func_is_aggregation():
+@pytest.mark.parametrize("chunked_array_api", CHUNKED_ARRAY_BACKENDS, indirect=True)
+def test_func_is_aggregation(chunked_array_api):
     from flox.aggregations import mean
 
     ds = xr.Dataset(
         {
             "Tair": (
                 ("time", "x", "y"),
-                dask.array.ones((36, 205, 275), chunks=(9, -1, -1)),
+                chunked_array_api.ones((36, 205, 275), chunks=(9, -1, -1)),
             )
         },
         coords={"time": xr.date_range("1980-09-01 00:00", "1983-09-18 00:00", freq="ME", calendar="noleap")},
diff --git a/uv.lock b/uv.lock
index 698fb951..1e3782d6 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1028,6 +1028,7 @@ docs = [
     { name = "myst-parser" },
     { name = "numpydoc" },
     { name = "pyarrow" },
+    { name = "rich" },
     { name = "sparse" },
     { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
     { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
@@ -1124,6 +1125,7 @@ docs = [
     { name = "myst-nb" },
     { name = "myst-parser" },
     { name = "numpydoc" },
+    { name = "rich" },
     { name = "sparse" },
     { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" },
     { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
@@ -1314,6 +1316,7 @@ requires-dist = [
     { name = "packaging", specifier = ">=21.3" },
     { name = "pandas", specifier = ">=2.1" },
     { name = "pyarrow", marker = "extra == 'docs'" },
+    { name = "rich", marker = "extra == 'docs'" },
     { name = "scipy", specifier = ">=1.12" },
     { name = "sparse", marker = "extra == 'docs'" },
     { name = "sphinx", marker = "extra == 'docs'" },
@@ -1411,6 +1414,7 @@ docs = [
     { name = "myst-nb" },
     { name = "myst-parser" },
     { name = "numpydoc" },
+    { name = "rich" },
     { name = "sparse" },
     { name = "sphinx" },
     { name = "sphinx-codeautolink" },