Repository: scverse/fast-array-utils
Pull request #46: Add mean_var
Author: login: flying-sheep
Created: 2025-03-03T12:09:27Z | Merged: 2025-03-04T16:16:50Z

Fixes #45 

Linked issue #45: New function: mean_var
https://github.com/scverse/scanpy/blob/9741ca674247f848162205a56ae62ddb1c670b17/src/scanpy/preprocessing/_utils.py#L34-L152

Diff:
diff --git a/src/fast_array_utils/stats/__init__.py b/src/fast_array_utils/stats/__init__.py
index 64a7b5a..4a5e01a 100644
--- a/src/fast_array_utils/stats/__init__.py
+++ b/src/fast_array_utils/stats/__init__.py
@@ -5,7 +5,8 @@
 
 from ._is_constant import is_constant
 from ._mean import mean
+from ._mean_var import mean_var
 from ._sum import sum
 
 
-__all__ = ["is_constant", "mean", "sum"]
+__all__ = ["is_constant", "mean", "mean_var", "sum"]
diff --git a/src/fast_array_utils/stats/_mean.py b/src/fast_array_utils/stats/_mean.py
index ec7dfe2..92c7345 100644
--- a/src/fast_array_utils/stats/_mean.py
+++ b/src/fast_array_utils/stats/_mean.py
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: MPL-2.0
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, overload
 
 import numpy as np
 
@@ -13,19 +13,34 @@
 
     from numpy._typing._array_like import _ArrayLikeFloat_co as ArrayLike
     from numpy.typing import DTypeLike, NDArray
+    from optype.numpy import ToDType
 
     from .. import types
 
-    # all supported types except OutOfCoreDataset (TODO)
-    Array = (
+    # all supported types except Dask and OutOfCoreDataset (TODO)
+    NonDaskArray = (
         NDArray[Any]
         | types.CSBase
         | types.H5Dataset
         | types.ZarrArray
         | types.CupyArray
         | types.CupySparseMatrix
-        | types.DaskArray
     )
+    Array = NonDaskArray | types.DaskArray
+
+
+@overload
+def mean(
+    x: ArrayLike | NonDaskArray, /, *, axis: Literal[None] = None, dtype: DTypeLike | None = None
+) -> np.number[Any]: ...
+@overload
+def mean(
+    x: ArrayLike | NonDaskArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None
+) -> NDArray[np.number[Any]]: ...
+@overload
+def mean(
+    x: types.DaskArray, /, *, axis: Literal[0, 1], dtype: ToDType[Any] | None = None
+) -> types.DaskArray: ...
 
 
 def mean(
@@ -34,7 +49,7 @@ def mean(
     *,
     axis: Literal[0, 1, None] = None,
     dtype: DTypeLike | None = None,
-) -> NDArray[Any] | types.DaskArray:
+) -> NDArray[np.number[Any]] | np.number[Any] | types.DaskArray:
     """Mean over both or one axis.
 
     Returns
@@ -49,7 +64,7 @@ def mean(
     if not hasattr(x, "shape"):
         raise NotImplementedError  # TODO(flying-sheep): infer shape  # noqa: TD003
     if TYPE_CHECKING:
-        assert isinstance(x, Array)
+        assert isinstance(x, Array)  # type:ignore[unused-ignore]
     total = sum_(x, axis=axis, dtype=dtype)
     n = np.prod(x.shape) if axis is None else x.shape[axis]
     return total / n
diff --git a/src/fast_array_utils/stats/_mean_var.py b/src/fast_array_utils/stats/_mean_var.py
new file mode 100644
index 0000000..e1d65fa
--- /dev/null
+++ b/src/fast_array_utils/stats/_mean_var.py
@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: MPL-2.0
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, no_type_check, overload
+
+import numba
+import numpy as np
+
+from .. import types
+from ._mean import mean
+from ._power import power
+
+
+if TYPE_CHECKING:
+    from typing import Any, Literal
+
+    from numpy.typing import NDArray
+
+    MemArray = NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix
+
+
+__all__ = ["mean_var"]
+
+
+@overload
+def mean_var(
+    x: MemArray, /, *, axis: Literal[None] = None, correction: int = 0
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]: ...
+@overload
+def mean_var(
+    x: MemArray, /, *, axis: Literal[0, 1], correction: int = 0
+) -> tuple[np.float64, np.float64]: ...
+@overload
+def mean_var(
+    x: types.DaskArray, /, *, axis: Literal[0, 1, None] = None, correction: int = 0
+) -> tuple[types.DaskArray, types.DaskArray]: ...
+
+
+@no_type_check  # mypy is extremely confused
+def mean_var(
+    x: MemArray | types.DaskArray,
+    /,
+    *,
+    axis: Literal[0, 1, None] = None,
+    correction: int = 0,
+) -> (
+    tuple[NDArray[np.float64], NDArray[np.float64]]
+    | tuple[np.float64, np.float64]
+    | tuple[types.DaskArray, types.DaskArray]
+):
+    if axis is not None and isinstance(x, types.CSBase):
+        mean_, var = _sparse_mean_var(x, axis=axis)
+    else:
+        mean_ = mean(x, axis=axis, dtype=np.float64)
+        mean_sq = mean(power(x, 2), axis=axis, dtype=np.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]
+        if n != 1:
+            var *= n / (n - correction)
+    return mean_, var
+
+
+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.
+
+    This code and internal functions are based on sklearns `sparsefuncs.mean_variance_axis`.
+
+    Modifications:
+    - allow deciding on the output type,
+      which can increase accuracy when calculating the mean and variance of 32bit floats.
+    - Doesn't currently implement support for null values, but could.
+    - Uses numba instead of Cython
+    """
+    assert axis in (0, 1)
+    if mtx.format == "csr":
+        ax_minor = 1
+        shape = mtx.shape
+    elif mtx.format == "csc":
+        ax_minor = 0
+        shape = mtx.shape[::-1]
+    else:
+        msg = "This function only works on sparse csr and csc matrices"
+        raise TypeError(msg)
+    if len(shape) == 1:
+        msg = "array must have 2 dimensions"
+        raise TypeError(msg)
+    f = sparse_mean_var_major_axis if axis == ax_minor else sparse_mean_var_minor_axis
+    return f(
+        mtx.data,
+        mtx.indptr,
+        mtx.indices,
+        major_len=shape[0],
+        minor_len=shape[1],
+        n_threads=numba.get_num_threads(),
+    )
+
+
+@numba.njit
+def sparse_mean_var_minor_axis(
+    data: NDArray[np.number[Any]],
+    indptr: NDArray[np.integer[Any]],
+    indices: NDArray[np.integer[Any]],
+    *,
+    major_len: int,
+    minor_len: int,
+    n_threads: int,
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    """Compute mean and variance along the minor axis of a compressed sparse matrix."""
+    rows = len(indptr) - 1
+    sums = np.zeros((n_threads, minor_len))
+    squared_sums = np.zeros((n_threads, minor_len))
+    means = np.zeros(minor_len)
+    variances = np.zeros(minor_len)
+    for i in numba.prange(n_threads):
+        for r in range(i, rows, n_threads):
+            for j in range(indptr[r], indptr[r + 1]):
+                minor_index = indices[j]
+                if minor_index >= minor_len:
+                    continue
+                value = data[j]
+                sums[i, minor_index] += value
+                squared_sums[i, minor_index] += value * value
+    for c in numba.prange(minor_len):
+        sum = sums[:, c].sum()
+        means[c] = sum / major_len
+        variances[c] = squared_sums[:, c].sum() / major_len - (sum / major_len) ** 2
+    return means, variances
+
+
+@numba.njit
+def sparse_mean_var_major_axis(
+    data: NDArray[np.number[Any]],
+    indptr: NDArray[np.integer[Any]],
+    indices: NDArray[np.integer[Any]],  # noqa: ARG001
+    *,
+    major_len: int,
+    minor_len: int,
+    n_threads: int,
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    """Compute means and variances along the major axis of a compressed sparse matrix."""
+    rows = len(indptr) - 1
+    means = np.zeros(major_len)
+    variances = np.zeros_like(means)
+
+    for i in numba.prange(n_threads):
+        for r in range(i, rows, n_threads):
+            sum_major = np.float64(0.0)
+            squared_sum_minor = np.float64(0.0)
+            for j in range(indptr[r], indptr[r + 1]):
+                value = np.float64(data[j])
+                sum_major += value
+                squared_sum_minor += value * value
+            means[r] = sum_major
+            variances[r] = squared_sum_minor
+    for c in numba.prange(major_len):
+        mean = means[c] / minor_len
+        means[c] = mean
+        variances[c] = variances[c] / minor_len - mean * mean
+    return means, variances
diff --git a/src/fast_array_utils/stats/_power.py b/src/fast_array_utils/stats/_power.py
new file mode 100644
index 0000000..3e15d0c
--- /dev/null
+++ b/src/fast_array_utils/stats/_power.py
@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: MPL-2.0
+from __future__ import annotations
+
+from functools import singledispatch
+from typing import TYPE_CHECKING, cast
+
+from .. import types
+
+
+if TYPE_CHECKING:
+    from typing import Any, TypeVar
+
+    from numpy.typing import NDArray
+
+    # All supported array types except for disk ones and OutOfCoreDataset
+    Array = NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix | types.DaskArray
+
+    _Arr = TypeVar("_Arr", bound=Array)
+
+
+def power(x: _Arr, n: int, /) -> _Arr:
+    """Take array or matrix to a power."""
+    # This wrapper is necessary because TypeVars can’t be used in `singledispatch` functions
+    return _power(x, n)  # type: ignore[return-value]
+
+
+@singledispatch
+def _power(x: Array, n: int, /) -> Array:
+    if TYPE_CHECKING:
+        assert not isinstance(x, types.DaskArray | types.CSMatrix)
+    return x**n  # type: ignore[operator]
+
+
+@_power.register(types.CSMatrix)  # type: ignore[call-overload,misc]
+def _power_cs(x: types.CSMatrix, n: int, /) -> types.CSMatrix:
+    return x.power(n)
+
+
+@_power.register(types.DaskArray)
+def _power_dask(x: types.DaskArray, n: int, /) -> types.DaskArray:
+    import dask.array as da
+
+    return cast(types.DaskArray, da.map_blocks(power, x, n))  # type: ignore[no-untyped-call]
diff --git a/stubs/cupy.pyi b/stubs/cupy.pyi
index ff2c6a9..8664e9e 100644
--- a/stubs/cupy.pyi
+++ b/stubs/cupy.pyi
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: MPL-2.0
-from typing import Any, Literal
+from typing import Any, Literal, Self
 
 import numpy as np
 from numpy.typing import ArrayLike, DTypeLike, NDArray
@@ -8,6 +8,7 @@ class ndarray:
     dtype: np.dtype[Any]
     shape: tuple[int, ...]
     def get(self) -> NDArray[Any]: ...
+    def __power__(self, other: int) -> Self: ...
 
 def asarray(
     a: ArrayLike,
diff --git a/stubs/cupyx/scipy/sparse.pyi b/stubs/cupyx/scipy/sparse.pyi
index 80973d6..f8011f5 100644
--- a/stubs/cupyx/scipy/sparse.pyi
+++ b/stubs/cupyx/scipy/sparse.pyi
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: MPL-2.0
-from typing import Any, Literal
+from typing import Any, Literal, Self
 
 import cupy
 import numpy as np
@@ -8,3 +8,4 @@ class spmatrix:
     dtype: np.dtype[Any]
     shape: tuple[int, int]
     def toarray(self, order: Literal["C", "F", None] = None, out: None = None) -> cupy.ndarray: ...
+    def __power__(self, other: int) -> Self: ...
diff --git a/stubs/numba/__init__.pyi b/stubs/numba/__init__.pyi
index 4785482..11b164d 100644
--- a/stubs/numba/__init__.pyi
+++ b/stubs/numba/__init__.pyi
@@ -31,3 +31,4 @@ def prange(stop: SupportsIndex, /) -> Iterable[int]: ...
 def prange(
     start: SupportsIndex, stop: SupportsIndex, step: SupportsIndex = ..., /
 ) -> Iterable[int]: ...
+def get_num_threads() -> int: ...
diff --git a/tests/test_stats.py b/tests/test_stats.py
index fb5e76c..9484c57 100644
--- a/tests/test_stats.py
+++ b/tests/test_stats.py
@@ -90,15 +90,41 @@ def test_sum(
 
 @pytest.mark.parametrize(("axis", "expected"), [(None, 3.5), (0, [2.5, 3.5, 4.5]), (1, [2.0, 5.0])])
 def test_mean(
-    array_type: ArrayType[Array], axis: Literal[0, 1, None], expected: list[float]
+    array_type: ArrayType[Array], axis: Literal[0, 1, None], expected: float | list[float]
 ) -> None:
-    arr = array_type(np.array([[1, 2, 3], [4, 5, 6]]))
+    np_arr = np.array([[1, 2, 3], [4, 5, 6]])
+    np.testing.assert_array_equal(np.mean(np_arr, axis=axis), expected)
+
+    arr = array_type(np_arr)
     result = stats.mean(arr, axis=axis)
     if isinstance(result, types.DaskArray):
-        result = result.compute()  # type: ignore[no-untyped-call]
+        result = result.compute()
     np.testing.assert_array_equal(result, expected)
 
 
+@pytest.mark.array_type(skip=Flags.Disk)
+@pytest.mark.parametrize(
+    ("axis", "mean_expected", "var_expected"),
+    [(None, 3.5, 3.5), (0, [2.5, 3.5, 4.5], [4.5, 4.5, 4.5]), (1, [2.0, 5.0], [1.0, 1.0])],
+)
+def test_mean_var(
+    array_type: ArrayType[
+        NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix | types.DaskArray
+    ],
+    axis: Literal[0, 1, None],
+    mean_expected: float | list[float],
+    var_expected: float | list[float],
+) -> None:
+    np_arr = np.array([[1, 2, 3], [4, 5, 6]])
+    np.testing.assert_array_equal(np.mean(np_arr, axis=axis), mean_expected)
+    np.testing.assert_array_equal(np.var(np_arr, axis=axis, correction=1), var_expected)
+
+    arr = array_type(np_arr)
+    mean, var = stats.mean_var(arr, axis=axis, correction=1)
+    np.testing.assert_array_equal(mean, mean_expected)
+    np.testing.assert_array_almost_equal_nulp(var, var_expected, nulp=8)
+
+
 # TODO(flying-sheep): enable for GPU  # noqa: TD003
 @pytest.mark.array_type(skip=Flags.Disk | Flags.Gpu)
 @pytest.mark.parametrize(
@@ -149,7 +175,7 @@ def test_dask_constant_blocks(
 
 @pytest.mark.benchmark
 @pytest.mark.array_type(skip=Flags.Matrix | Flags.Dask | Flags.Disk | Flags.Gpu)
-@pytest.mark.parametrize("func", [stats.sum, stats.mean, stats.is_constant])
+@pytest.mark.parametrize("func", [stats.sum, stats.mean, stats.mean_var, stats.is_constant])
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])  # random only supports float
 def test_stats_benchmark(
     benchmark: BenchmarkFixture,


Reviews:
user: login: ilan-gold
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-04T12:48:36Z
user: login: flying-sheep
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-04T13:51:11Z
user: login: flying-sheep
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-04T13:55:28Z
user: login: ilan-gold
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-04T15:13:17Z
user: login: ilan-gold
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-04T15:16:49Z
user: login: ilan-gold
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-04T15:16:54Z

Review comments:
diff_hunk: @@ -0,0 +1,42 @@
+# SPDX-License-Identifier: MPL-2.0
+from __future__ import annotations
+
+from functools import singledispatch
+from typing import TYPE_CHECKING, cast
+
+from .. import types
+
+
+if TYPE_CHECKING:
+    from typing import Any, TypeVar
+
+    from numpy.typing import NDArray
+
+    # All supported array types except for disk ones and OutOfCoreDataset
+    Array = NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix | types.DaskArray
+
+    _Arr = TypeVar("_Arr", bound=Array)
+
+
+def power(x: _Arr, n: int, /) -> _Arr:
+    """Take array or matrix to a power."""
+    return _power(x, n)  # type: ignore[return-value]
path: src/fast_array_utils/stats/_power.py
user: login: ilan-gold
body: Why is this wrapper here?
created_at: 2025-03-04T12:44:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -90,15 +90,41 @@ def test_sum(
 
 @pytest.mark.parametrize(("axis", "expected"), [(None, 3.5), (0, [2.5, 3.5, 4.5]), (1, [2.0, 5.0])])
 def test_mean(
-    array_type: ArrayType[Array], axis: Literal[0, 1, None], expected: list[float]
+    array_type: ArrayType[Array], axis: Literal[0, 1, None], expected: float | list[float]
 ) -> None:
-    arr = array_type(np.array([[1, 2, 3], [4, 5, 6]]))
+    np_arr = np.array([[1, 2, 3], [4, 5, 6]])
+    np.testing.assert_array_equal(np.mean(np_arr, axis=axis), expected)
+
+    arr = array_type(np_arr)
     result = stats.mean(arr, axis=axis)
     if isinstance(result, types.DaskArray):
-        result = result.compute()  # type: ignore[no-untyped-call]
+        result = result.compute()
     np.testing.assert_array_equal(result, expected)
 
 
+@pytest.mark.array_type(skip=Flags.Disk)
+@pytest.mark.parametrize(
+    ("axis", "mean_expected", "var_expected"),
+    [(None, 3.5, 3.5), (0, [2.5, 3.5, 4.5], [4.5, 4.5, 4.5]), (1, [2.0, 5.0], [1.0, 1.0])],
+)
+def test_mean_var(
+    array_type: ArrayType[
+        NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix | types.DaskArray
+    ],
+    axis: Literal[0, 1, None],
+    mean_expected: float | list[float],
+    var_expected: float | list[float],
+) -> None:
+    np_arr = np.array([[1, 2, 3], [4, 5, 6]])
+    np.testing.assert_array_equal(np.mean(np_arr, axis=axis), mean_expected)
+    np.testing.assert_array_equal(np.var(np_arr, axis=axis, correction=1), var_expected)
+
+    arr = array_type(np_arr)
+    mean, var = stats.mean_var(arr, axis=axis, correction=1)
+    np.testing.assert_array_equal(mean, mean_expected)
+    np.testing.assert_array_almost_equal(var, var_expected)
path: tests/test_stats.py
user: login: ilan-gold
body: See note: https://numpy.org/doc/2.2/reference/generated/numpy.testing.assert_array_almost_equal.html , should we switch to [`assert_allclose`](https://numpy.org/doc/2.2/reference/generated/numpy.testing.assert_allclose.html#numpy.testing.assert_allclose)? 
created_at: 2025-03-04T12:48:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,42 @@
+# SPDX-License-Identifier: MPL-2.0
+from __future__ import annotations
+
+from functools import singledispatch
+from typing import TYPE_CHECKING, cast
+
+from .. import types
+
+
+if TYPE_CHECKING:
+    from typing import Any, TypeVar
+
+    from numpy.typing import NDArray
+
+    # All supported array types except for disk ones and OutOfCoreDataset
+    Array = NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix | types.DaskArray
+
+    _Arr = TypeVar("_Arr", bound=Array)
+
+
+def power(x: _Arr, n: int, /) -> _Arr:
+    """Take array or matrix to a power."""
+    return _power(x, n)  # type: ignore[return-value]
path: src/fast_array_utils/stats/_power.py
user: login: flying-sheep
body: because of the stupid way `singledispatch`’s typing works (with the fallback function’s first argument having to be the supertype of all registered functions’ first arguments):

```
src/fast_array_utils/stats/_power.py:35: error: Dispatch type "Array" must be subtype of fallback function first argument "_Arr"  [misc]
```

That means
1. when the return value is in any way dependent on the input, one has to wrap a `singledispatch`ed function to get correct typing.
2. the first argument’s type in the fallback is incorrect, thus the `if TYPE_CHECKING: assert isinstance`.

See also #27
created_at: 2025-03-04T13:51:11Z
author_association: MEMBER
diff_hunk: @@ -90,15 +90,41 @@ def test_sum(
 
 @pytest.mark.parametrize(("axis", "expected"), [(None, 3.5), (0, [2.5, 3.5, 4.5]), (1, [2.0, 5.0])])
 def test_mean(
-    array_type: ArrayType[Array], axis: Literal[0, 1, None], expected: list[float]
+    array_type: ArrayType[Array], axis: Literal[0, 1, None], expected: float | list[float]
 ) -> None:
-    arr = array_type(np.array([[1, 2, 3], [4, 5, 6]]))
+    np_arr = np.array([[1, 2, 3], [4, 5, 6]])
+    np.testing.assert_array_equal(np.mean(np_arr, axis=axis), expected)
+
+    arr = array_type(np_arr)
     result = stats.mean(arr, axis=axis)
     if isinstance(result, types.DaskArray):
-        result = result.compute()  # type: ignore[no-untyped-call]
+        result = result.compute()
     np.testing.assert_array_equal(result, expected)
 
 
+@pytest.mark.array_type(skip=Flags.Disk)
+@pytest.mark.parametrize(
+    ("axis", "mean_expected", "var_expected"),
+    [(None, 3.5, 3.5), (0, [2.5, 3.5, 4.5], [4.5, 4.5, 4.5]), (1, [2.0, 5.0], [1.0, 1.0])],
+)
+def test_mean_var(
+    array_type: ArrayType[
+        NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix | types.DaskArray
+    ],
+    axis: Literal[0, 1, None],
+    mean_expected: float | list[float],
+    var_expected: float | list[float],
+) -> None:
+    np_arr = np.array([[1, 2, 3], [4, 5, 6]])
+    np.testing.assert_array_equal(np.mean(np_arr, axis=axis), mean_expected)
+    np.testing.assert_array_equal(np.var(np_arr, axis=axis, correction=1), var_expected)
+
+    arr = array_type(np_arr)
+    mean, var = stats.mean_var(arr, axis=axis, correction=1)
+    np.testing.assert_array_equal(mean, mean_expected)
+    np.testing.assert_array_almost_equal(var, var_expected)
path: tests/test_stats.py
user: login: flying-sheep
body: I don’t understand rtol and atol. Maybe `assert_array_almost_equal_nulp` is a better replacement. would be:

```suggestion
    np.testing.assert_array_almost_equal_nulp(var, var_expected, nulp=8)
```
created_at: 2025-03-04T13:55:28Z
author_association: MEMBER
diff_hunk: @@ -90,15 +90,41 @@ def test_sum(
 
 @pytest.mark.parametrize(("axis", "expected"), [(None, 3.5), (0, [2.5, 3.5, 4.5]), (1, [2.0, 5.0])])
 def test_mean(
-    array_type: ArrayType[Array], axis: Literal[0, 1, None], expected: list[float]
+    array_type: ArrayType[Array], axis: Literal[0, 1, None], expected: float | list[float]
 ) -> None:
-    arr = array_type(np.array([[1, 2, 3], [4, 5, 6]]))
+    np_arr = np.array([[1, 2, 3], [4, 5, 6]])
+    np.testing.assert_array_equal(np.mean(np_arr, axis=axis), expected)
+
+    arr = array_type(np_arr)
     result = stats.mean(arr, axis=axis)
     if isinstance(result, types.DaskArray):
-        result = result.compute()  # type: ignore[no-untyped-call]
+        result = result.compute()
     np.testing.assert_array_equal(result, expected)
 
 
+@pytest.mark.array_type(skip=Flags.Disk)
+@pytest.mark.parametrize(
+    ("axis", "mean_expected", "var_expected"),
+    [(None, 3.5, 3.5), (0, [2.5, 3.5, 4.5], [4.5, 4.5, 4.5]), (1, [2.0, 5.0], [1.0, 1.0])],
+)
+def test_mean_var(
+    array_type: ArrayType[
+        NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix | types.DaskArray
+    ],
+    axis: Literal[0, 1, None],
+    mean_expected: float | list[float],
+    var_expected: float | list[float],
+) -> None:
+    np_arr = np.array([[1, 2, 3], [4, 5, 6]])
+    np.testing.assert_array_equal(np.mean(np_arr, axis=axis), mean_expected)
+    np.testing.assert_array_equal(np.var(np_arr, axis=axis, correction=1), var_expected)
+
+    arr = array_type(np_arr)
+    mean, var = stats.mean_var(arr, axis=axis, correction=1)
+    np.testing.assert_array_equal(mean, mean_expected)
+    np.testing.assert_array_almost_equal(var, var_expected)
path: tests/test_stats.py
user: login: ilan-gold
body: yup fine by me
created_at: 2025-03-04T15:13:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,42 @@
+# SPDX-License-Identifier: MPL-2.0
+from __future__ import annotations
+
+from functools import singledispatch
+from typing import TYPE_CHECKING, cast
+
+from .. import types
+
+
+if TYPE_CHECKING:
+    from typing import Any, TypeVar
+
+    from numpy.typing import NDArray
+
+    # All supported array types except for disk ones and OutOfCoreDataset
+    Array = NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix | types.DaskArray
+
+    _Arr = TypeVar("_Arr", bound=Array)
+
+
+def power(x: _Arr, n: int, /) -> _Arr:
+    """Take array or matrix to a power."""
+    return _power(x, n)  # type: ignore[return-value]
path: src/fast_array_utils/stats/_power.py
user: login: ilan-gold
body: Ok I would recommend adding a small note here then
created_at: 2025-03-04T15:16:49Z
author_association: CONTRIBUTOR

PR comments:
user: login: codecov[bot]
created_at: 2025-03-03T12:10:14Z
body: ## [Codecov](https://app.codecov.io/gh/scverse/fast-array-utils/pull/46?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse) Report
Attention: Patch coverage is `92.59259%` with `4 lines` in your changes missing coverage. Please review.
> Project coverage is 94.44%. Comparing base [(`a38413d`)](https://app.codecov.io/gh/scverse/fast-array-utils/commit/a38413da89a7f61fcdaad0c676c76562b445e5f7?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse) to head [(`8317937`)](https://app.codecov.io/gh/scverse/fast-array-utils/commit/8317937bbfe4bb25519b593ac778ce2eff8cd048?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse).
> Report is 2 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/scverse/fast-array-utils/pull/46?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse) | Patch % | Lines |
|---|---|---|
| [src/fast\_array\_utils/stats/\_mean\_var.py](https://app.codecov.io/gh/scverse/fast-array-utils/pull/46?src=pr&el=tree&filepath=src%2Ffast_array_utils%2Fstats%2F_mean_var.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c3JjL2Zhc3RfYXJyYXlfdXRpbHMvc3RhdHMvX21lYW5fdmFyLnB5) | 88.57% | [4 Missing :warning: ](https://app.codecov.io/gh/scverse/fast-array-utils/pull/46?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main      #46      +/-   ##
==========================================
- Coverage   95.05%   94.44%   -0.62%     
==========================================
  Files          12       14       +2     
  Lines         182      234      +52     
==========================================
+ Hits          173      221      +48     
- Misses          9       13       +4     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/scverse/fast-array-utils/pull/46?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse).

author_association: NONE
user: login: codspeed-hq[bot]
created_at: 2025-03-03T12:12:28Z
body: <!-- __CODSPEED_PERFORMANCE_REPORT_COMMENT__ -->
<!-- __CODSPEED_INSTRUMENTATION_PERFORMANCE_REPORT_COMMENT__ -->

## [CodSpeed Performance Report](https://codspeed.io/scverse/fast-array-utils/branches/pa%2Fmean-var)

### Merging #46 will **degrade performances by 52.33%**

<sub>Comparing <code>pa/mean-var</code> (8317937) with <code>main</code> (2c7a94a)</sub>



### Summary

`⚡ 1` improvements  
`❌ 1` regressions  
`✅ 68` untouched benchmarks  
`🆕 18` new benchmarks  


> :warning: _Please fix the performance issues or [acknowledge them on CodSpeed](https://codspeed.io/scverse/fast-array-utils/branches/pa%2Fmean-var)._

### Benchmarks breakdown

|     | Benchmark | `BASE` | `HEAD` | Change |
| --- | --------- | ----------------------- | ------------------- | ------ |
| 🆕 | `` test_stats_benchmark[0-numpy.ndarray-float32-mean_var] `` | N/A | 345.7 µs | N/A |
| 🆕 | `` test_stats_benchmark[0-numpy.ndarray-float64-mean_var] `` | N/A | 326.2 µs | N/A |
| 🆕 | `` test_stats_benchmark[0-scipy.sparse.csc_array-float32-mean_var] `` | N/A | 164.1 µs | N/A |
| 🆕 | `` test_stats_benchmark[0-scipy.sparse.csc_array-float64-mean_var] `` | N/A | 180.5 µs | N/A |
| 🆕 | `` test_stats_benchmark[0-scipy.sparse.csr_array-float32-mean_var] `` | N/A | 335.5 µs | N/A |
| 🆕 | `` test_stats_benchmark[0-scipy.sparse.csr_array-float64-mean_var] `` | N/A | 348.8 µs | N/A |
| 🆕 | `` test_stats_benchmark[1-numpy.ndarray-float32-mean_var] `` | N/A | 339.1 µs | N/A |
| 🆕 | `` test_stats_benchmark[1-numpy.ndarray-float64-mean_var] `` | N/A | 320.9 µs | N/A |
| 🆕 | `` test_stats_benchmark[1-scipy.sparse.csc_array-float32-mean_var] `` | N/A | 338.8 µs | N/A |
| 🆕 | `` test_stats_benchmark[1-scipy.sparse.csc_array-float64-mean_var] `` | N/A | 353 µs | N/A |
| ❌ | `` test_stats_benchmark[1-scipy.sparse.csr_array-float32-is_constant] `` | 97.6 µs | 204.7 µs | -52.33% |
| 🆕 | `` test_stats_benchmark[1-scipy.sparse.csr_array-float32-mean_var] `` | N/A | 160.7 µs | N/A |
| ⚡ | `` test_stats_benchmark[1-scipy.sparse.csr_array-float64-is_constant] `` | 216.3 µs | 96.1 µs | ×2.3 |
| 🆕 | `` test_stats_benchmark[1-scipy.sparse.csr_array-float64-mean_var] `` | N/A | 176.6 µs | N/A |
| 🆕 | `` test_stats_benchmark[None-numpy.ndarray-float32-mean_var] `` | N/A | 346.3 µs | N/A |
| 🆕 | `` test_stats_benchmark[None-numpy.ndarray-float64-mean_var] `` | N/A | 331 µs | N/A |
| 🆕 | `` test_stats_benchmark[None-scipy.sparse.csc_array-float32-mean_var] `` | N/A | 698.6 µs | N/A |
| 🆕 | `` test_stats_benchmark[None-scipy.sparse.csc_array-float64-mean_var] `` | N/A | 740.2 µs | N/A |
| 🆕 | `` test_stats_benchmark[None-scipy.sparse.csr_array-float32-mean_var] `` | N/A | 679.5 µs | N/A |
| 🆕 | `` test_stats_benchmark[None-scipy.sparse.csr_array-float64-mean_var] `` | N/A | 714.7 µs | N/A |

author_association: NONE