fix: chan's parallel mean-var algorithm for dask-backed arrays

.github/workflows/ci.yml

@@ -54,16 +54,18 @@ jobs:
     steps:
       - uses: actions/checkout@v7
         with: { fetch-depth: 0, filter: "blob:none" }
+      # https://github.com/pypa/hatch/pull/2382
+      - run: uv tool install 'hatch @ git+https://github.com/pypa/hatch.git@ed8e30bebf98f2fe4d70c18a32a50a8160c391cb'
       - uses: astral-sh/setup-uv@v9.0.0
         with:
           python-version: ${{ matrix.env.python }}
       - name: create environment
-        run: uvx hatch -v env create ${{ matrix.env.name }}
+        run: hatch -v env create ${{ matrix.env.name }}
       - name: run tests with coverage
         run: |
-          uvx hatch run ${{ matrix.env.name }}:run-cov
+          hatch run ${{ matrix.env.name }}:run-cov
           # https://github.com/codecov/codecov-cli/issues/648
-          uvx hatch run ${{ matrix.env.name }}:coverage xml
+          hatch run ${{ matrix.env.name }}:coverage xml
           rm test-data/.coverage
       - uses: codecov/codecov-action@v7
         with:

.vscode/settings.json

@@ -1,10 +1,7 @@
 {
-    "[toml][json][jsonc][python]": {
+    "[json][jsonc][python]": {
         "editor.formatOnSave": true,
     },
-    "[toml]": {
-        "editor.defaultFormatter": "tamasfe.even-better-toml",
-    },
     "[json][jsonc]": {
         "editor.defaultFormatter": "biomejs.biome",
     },

pyproject.toml

@@ -25,7 +25,7 @@ classifiers = [
 dynamic = [ "version" ]
 dependencies = [ "array-api-compat", "numpy>=2" ]
 optional-dependencies.accel = [ "numba>=0.57" ]
-optional-dependencies.dask = [ "dask>=2023.6.1" ]
+optional-dependencies.dask = [ "dask>=2024.5.1" ]
 optional-dependencies.full = [ "fast-array-utils[accel,dask,sparse]", "h5py", "zarr" ]
 optional-dependencies.sparse = [ "scipy>=1.13" ]
 optional-dependencies.testing = [ "packaging" ]
@@ -86,7 +86,7 @@ envs.hatch-test.overrides.matrix.extras.dependency-groups = [
 ]
 envs.hatch-test.overrides.matrix.resolution.dependencies = [
   { value = "numpy==2", if = [ "lowest" ] },
-  { value = "dask==2023.6.1", if = [ "lowest" ] },
+  { value = "dask==2024.5.1", if = [ "lowest" ] },
   { value = "scipy==1.13.0", if = [ "lowest" ] },
 ]
 envs.hatch-test.default-args = []

src/fast_array_utils/stats/__init__.py

@@ -91,7 +91,7 @@ def mean(x: CpuArray | DiskArray, /, *, axis: Literal[0, 1], dtype: DTypeLike |
 @overload
 def mean(x: GpuArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None) -> types.CupyArray: ...
 @overload
-def mean(x: types.DaskArray, /, *, axis: Literal[0, 1], dtype: ToDType[Any] | None = None) -> types.DaskArray: ...
+def mean(x: types.DaskArray, /, *, axis: Literal[0, 1] | None = None, dtype: ToDType[Any] | None = None) -> types.DaskArray: ...
 @overload
 def mean[A: types.HasArrayNamespace](x: A, /, *, axis: Literal[0, 1] | None = None, dtype: DTypeLike | None = None) -> A: ...
 

src/fast_array_utils/stats/_mean_var.py

@@ -1,22 +1,41 @@
 # SPDX-License-Identifier: MPL-2.0
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, no_type_check
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, cast, no_type_check
 
 import numba
 import numpy as np
 
 from .. import types
 from ..numba import njit
 from ._power import power
+from ._utils import _get_shape, _normalize_axis
 
 
 if TYPE_CHECKING:
-    from typing import Literal
+    from collections.abc import Iterator
+    from typing import Any, Literal, TypedDict
 
+    from dask.array.reductions import _Chunk
     from numpy.typing import NDArray
 
     from ..typing import CpuArray, GpuArray
+    from ._utils import ComplexAxis
+
+    class _Moments(TypedDict):
+        """A (count, mean, M2) triple as tracked by Chan's parallel-variance algorithm.
+
+        `mean`/`m2` are shaped like `_get_shape`'s `keepdims=True` convention.
+        """
+
+        n: int
+        mean: NDArray[np.float64]
+        m2: NDArray[np.float64]
+
+    # what dask passes to `combine`/`aggregate`: our own chunk output, a (nested) list
+    # thereof (`concatenate=False`), or a plain array while it’s computing `meta`
+    type MomentsIn = _Moments | CpuArray | GpuArray | Sequence[Any]
 
 
 @no_type_check  # mypy is extremely confused
@@ -32,6 +51,9 @@ def mean_var_(
     | tuple[np.float64, np.float64]
     | tuple[types.DaskArray, types.DaskArray]
 ):
+    if isinstance(x, types.DaskArray):
+        return _dask_mean_var(x, axis=axis, correction=correction)
+
     from . import mean
 
     if isinstance(x, np.ndarray | types.CSBase) or not isinstance(x, types.HasArrayNamespace):
@@ -45,7 +67,7 @@ def mean_var_(
         mean_, var = _sparse_mean_var(x, axis=axis)
     else:
         mean_ = mean(x, axis=axis, dtype=xp.float64)
-        mean_sq = mean(power(x, 2, dtype=xp.float64), axis=axis) if isinstance(x, types.DaskArray) else mean(power(x, 2), axis=axis, dtype=xp.float64)
+        mean_sq = mean(power(x, 2, dtype=xp.float64), axis=axis, dtype=xp.float64)
         var = mean_sq - mean_**2
     if correction:  # R convention == 1 (unbiased estimator)
         n = np.prod(x.shape) if axis is None else x.shape[axis]
@@ -54,6 +76,125 @@ def mean_var_(
     return mean_, var
 
 
+def _dask_mean_var(x: types.DaskArray, /, *, axis: Literal[0, 1] | None, correction: int) -> tuple[types.DaskArray, types.DaskArray]:
+    """Mean and variance of a dask array.
+
+    ``mean`` is a normal (associative) sum-based dask reduction.
+    ``var`` is instead derived from per-chunk ``(count, mean, M2)`` triples
+    (``M2 = sum((x - mean)**2)``, computed per chunk by recursing into :func:`mean_var_`)
+    that get merged pairwise across chunks using Chan's parallel-variance algorithm.
+    Separately dask-reducing ``mean(x)``/``mean(x**2)`` and subtracting at the end
+    (the naive two-pass formula used before) loses precision once many chunks are
+    combined, especially for float32 data on GPUs.
+    https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm
+    """
+    import dask.array as da
+
+    from . import mean
+
+    n = np.prod(x.shape) if axis is None else x.shape[axis]
+    mean_ = mean(x, axis=axis, dtype=np.float64)
+    # mypy can’t infer `reduction`’s type parameter from the callbacks, so pin it here
+    chunk: _Chunk[_Moments] = _moments_chunk
+    m2 = da.reduction(
+        x, chunk, _moments_aggregate, axis=axis, combine=_moments_combine, concatenate=False, dtype=np.float64, meta=np.array([], dtype=np.float64)
+    )
+    if axis is None:  # match `mean`/`sum`’s convention of reducing to a true scalar
+        m2 = m2.map_blocks(lambda a: a.reshape(())[()], meta=m2.dtype.type(0))
+    denom = n - correction if correction and n != 1 else n
+    return mean_, m2 / denom
+
+
+def _moments_chunk(
+    a: CpuArray | GpuArray,
+    /,
+    *,
+    axis: ComplexAxis = None,
+    keepdims: bool = False,
+    computing_meta: bool = False,
+    **kwargs: object,  # noqa: ARG001  # `_Chunk`/`_CB` let dask pass arbitrary keywords
+) -> _Moments | NDArray[np.float64]:
+    if computing_meta:  # pragma: no cover
+        return np.ndarray((), dtype=np.float64)
+    axis_ = _normalize_axis(axis, a.ndim)
+    mean_, var_ = mean_var_(a, axis=axis_, correction=0)
+    n = int(np.prod(a.shape)) if axis_ is None else a.shape[axis_]
+    shape = _get_shape(mean_, axis=axis_, keepdims=keepdims)
+    moments: _Moments = {"n": n, "mean": mean_.reshape(shape), "m2": (var_ * n).reshape(shape)}
+    return moments
+
+
+def _moments_combine(
+    pairs: MomentsIn,
+    /,
+    *,
+    axis: ComplexAxis = None,  # noqa: ARG001
+    keepdims: bool = False,  # noqa: ARG001
+    computing_meta: bool = False,
+    **kwargs: object,  # noqa: ARG001  # `_Chunk`/`_CB` let dask pass arbitrary keywords
+) -> _Moments | NDArray[np.float64]:
+    if computing_meta:  # pragma: no cover
+        return np.ndarray((), dtype=np.float64)
+    return _combine_all(pairs)
+
+
+def _moments_aggregate(
+    pairs: MomentsIn,
+    /,
+    *,
+    axis: ComplexAxis = None,
+    keepdims: bool = False,
+    computing_meta: bool = False,
+    **kwargs: object,  # noqa: ARG001  # `_Chunk`/`_CB` let dask pass arbitrary keywords
+) -> NDArray[np.float64]:
+    if computing_meta:  # pragma: no cover
+        return np.ndarray((), dtype=np.float64)
+    m2 = _combine_all(pairs)["m2"]
+    axis_ = _normalize_axis(axis, 2)
+    return m2.reshape(_final_moments_shape(m2.size, axis=axis_, keepdims=keepdims))
+
+
+def _combine_all(pairs: MomentsIn) -> _Moments:
+    """Merge every moment triple in a (possibly nested) list of `_moments_chunk` outputs."""
+    combined = None
+    for pair in _flatten_moments(pairs):
+        combined = pair if combined is None else _chan_combine(combined, pair)
+    assert combined is not None
+    return combined
+
+
+def _flatten_moments(pairs: MomentsIn) -> Iterator[_Moments]:
+    match pairs:
+        case {"n": _, "mean": _, "m2": _}:
+            yield cast("_Moments", pairs)
+        case Sequence():  # `concatenate=False` means dask hands us nested lists
+            for pair in pairs:
+                yield from _flatten_moments(pair)
+        case _:  # pragma: no cover
+            msg = f"Unexpected moments input: {type(pairs)}"
+            raise TypeError(msg)
+
+
+def _chan_combine(a: _Moments, b: _Moments) -> _Moments:
+    """Pairwise-merge two ``(count, mean, M2)`` moment triples."""
+    n_a, mean_a, m2_a = a["n"], a["mean"], a["m2"]
+    n_b, mean_b, m2_b = b["n"], b["mean"], b["m2"]
+    n = n_a + n_b
+    delta = mean_b - mean_a
+    mean_ = mean_a + delta * (n_b / n)
+    m2 = m2_a + m2_b + delta**2 * (n_a * n_b / n)
+    return {"n": n, "mean": mean_, "m2": m2}
+
+
+def _final_moments_shape(size: int, *, axis: Literal[0, 1] | None, keepdims: bool) -> tuple[int, ...]:
+    """Shape for a fully-combined moment array, mirroring `_get_shape`'s convention."""
+    if axis is None:
+        return (1, 1) if keepdims else (1,)
+    if not keepdims:
+        return (size,)
+    return (1, size) if axis == 0 else (size, 1)
+
+
 def _sparse_mean_var(mtx: types.CSBase, /, *, axis: Literal[0, 1]) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
     """Calculate means and variances for each row or column of a sparse matrix.
 

src/fast_array_utils/stats/_utils.py

@@ -19,7 +19,11 @@
     from ..typing import CpuArray, GpuArray
     from ._typing import DTypeKw, Ops
 
-    type ComplexAxis = tuple[Literal[0], Literal[1]] | tuple[Literal[0, 1]] | Literal[0, 1] | None
+    # `axis` as dask passes it into reduction callbacks:
+    # normalized to a tuple of non-negative ints,
+    # or not passed at all (hence `None`) when dask only computes a callback’s `meta`.
+    # `Literal[0, 1]` is what `_normalize_axis` narrows that down to.
+    type ComplexAxis = Literal[0, 1] | tuple[int, ...] | None
 
 
 __all__ = ["_dask_inner"]
@@ -85,12 +89,15 @@ def _normalize_axis(axis: ComplexAxis, ndim: int) -> Literal[0, 1] | None:
     match axis:
         case int() | None:  # pragma: no cover
             pass
-        case (0 | 1,):
-            axis = axis[0]
+        case (0,):
+            axis = 0
+        case (1,):
+            axis = 1
         case (0, 1) | (1, 0):
             axis = None
         case _:  # pragma: no cover
-            raise AxisError(axis, ndim)
+            msg = f"axis {axis!r} invalid for {ndim}-dimensional array"
+            raise AxisError(msg)
     if axis == 0 and ndim == 1:
         return None  # dask’s aggregate doesn’t know we don’t accept `axis=0` for 1D arrays
     return axis

src/testing/fast_array_utils/_array_type.py

@@ -311,7 +311,7 @@ def _to_scipy_sparse(
             x = to_dense(x, to_cpu_memory=True)  # type: ignore[arg-type]  # doesn’t officially handle ArrayLike
 
         cls = cast("type[types.CSBase]", cls or self.cls)
-        return cls(x, dtype=dtype)  # type: ignore[arg-type,misc]
+        return cls(x, dtype=dtype)  # type: ignore[arg-type]
 
     def _to_cupy_array(self, x: ArrayLike | Array, /, *, dtype: DTypeLike | None = None) -> types.CupyArray:
         import cupy as cu

tests/test_stats.py

@@ -316,10 +316,7 @@ def test_mean_var_sparse_32(array_type: ArrayType[types.CSArray], subtests: pyte
 
 @pytest.mark.array_type({at for at in SUPPORTED_TYPES if at.flags & Flags.Sparse and at.flags & Flags.Dask})
 def test_mean_var_pbmc_dask(array_type: ArrayType[types.DaskArray], pbmc64k_reduced_raw: sps.csr_array[np.float32]) -> None:
-    """Test float32 precision for bigger data.
-
-    This test is flaky for sparse-in-dask for some reason.
-    """
+    """Test float32 precision for bigger data."""
     mat = pbmc64k_reduced_raw
     arr = array_type(mat)
 

typings/dask/array/core.pyi

@@ -39,6 +39,7 @@ class Array:
     def __eq__(self, value: object, /) -> Array: ...  # type: ignore[override]
     def __getitem__(self, index: object) -> Array: ...
     def all(self) -> Array: ...
+    def __truediv__(self, other: float | np.number[Any] | Array, /) -> Array: ...
 
     # dask methods and attrs
     _meta: _Array

typings/dask/array/reductions.pyi

@@ -7,25 +7,40 @@ from numpy.typing import ArrayLike, DTypeLike, NDArray
 
 from .core import Array, _Array
 
-class _Chunk(Protocol):
+# dask only passes `computing_meta` to callbacks that accept it, so it’s optional here.
+class _Chunk[T](Protocol):
     @overload
-    def __call__(self, x_chunk: _Array, /, *, weights_chunk: NDArray[Any] | None = None, axis: tuple[int, ...], keepdims: bool, **kwargs: object) -> _Array: ...
+    def __call__(
+        self,
+        x_chunk: _Array,
+        /,
+        *,
+        weights_chunk: NDArray[Any] | None = None,
+        axis: tuple[int, ...],
+        keepdims: bool,
+        computing_meta: bool = ...,
+        **kwargs: object,
+    ) -> _Array | T: ...
     @overload
-    def __call__(self, x_chunk: _Array, /, *, axis: tuple[int, ...], keepdims: bool, **kwargs: object) -> _Array: ...
+    def __call__(self, x_chunk: _Array, /, *, axis: tuple[int, ...], keepdims: bool, computing_meta: bool = ..., **kwargs: object) -> _Array | T: ...
 
-class _CB(Protocol):
-    def __call__(self, x_chunk: _Array, /, *, axis: tuple[int, ...], keepdims: bool, **kwargs: object) -> _Array: ...
+class _CB[T](Protocol):
+    # When `concatenate=False`, dask passes a (possibly nested) list of the previous
+    # step’s raw outputs instead of concatenating them into a single `_Array`.
+    def __call__(
+        self, x_chunk: _Array | T | Sequence[Any], /, *, axis: tuple[int, ...], keepdims: bool, computing_meta: bool = ..., **kwargs: object
+    ) -> _Array | T: ...
 
-def reduction(
+def reduction[T](
     x: Array,
-    chunk: _Chunk,
-    aggregate: _CB,
+    chunk: _Chunk[T],
+    aggregate: _CB[T],
     *,
     axis: int | Sequence[int] | None = None,
     keepdims: bool = False,
     dtype: DTypeLike | None = None,
     split_every: int | Mapping[int, int] | None = None,
-    combine: _CB | None = None,
+    combine: _CB[T] | None = None,
     name: str | None = None,
     out: Array | None = None,
     concatenate: bool = True,