ENH: add nanmean

docs/sphinx/api-statistical.md

@@ -8,6 +8,7 @@
 
     cov
     nanmax
+    nanmean
     nanmin
     nansum
 ```

src/array_api_extra/__init__.py

@@ -13,7 +13,7 @@
 from ._searching import searchsorted
 from ._set import isin, nunique, setdiff1d, union1d
 from ._sorting import argpartition, partition
-from ._statistical import cov, nanmax, nanmin, nansum
+from ._statistical import cov, nanmax, nanmean, nanmin, nansum
 
 __version__ = "0.11.2.dev0"
 
@@ -37,6 +37,7 @@
     "lazy_apply",
     "nan_to_num",
     "nanmax",
+    "nanmean",
     "nanmin",
     "nansum",
     "nunique",

src/array_api_extra/_agnostic/_statistical.py

@@ -8,7 +8,7 @@
 from .._lib._typing import Array, ArrayNamespace
 from . import _manipulation
 
-__all__ = ["cov", "nanmax", "nanmin", "nansum"]
+__all__ = ["cov", "nanmax", "nanmean", "nanmin", "nansum"]
 
 
 def cov(
@@ -154,3 +154,27 @@ def nansum(  # numpydoc ignore=PR01,RT01
     """See docstring in `array_api_extra._statistical`."""
     mask = xp.isnan(a)
     return xp.sum(xp.where(mask, xp.zeros_like(a), a), axis=axis)
+
+
+def nanmean(  # numpydoc ignore=PR01,RT01
+    a: Array,
+    /,
+    *,
+    axis: int | tuple[int, ...] | None,
+    xp: ArrayNamespace,
+) -> Array:
+    """See docstring in `array_api_extra._statistical`."""
+    mask = xp.isnan(a)
+    sum_ = nansum(a, axis=axis, xp=xp)
+    count = xp.sum(xp.where(mask, xp.zeros_like(a), xp.ones_like(a)), axis=axis)
+    safe_count = xp.astype(
+        xp.where(count == 0, xp.ones_like(count), count),
+        sum_.dtype,
+        copy=False,
+    )
+    result = sum_ / safe_count
+    return xp.where(
+        count == 0,
+        xp.full_like(result, xp.nan),
+        result,
+    )

src/array_api_extra/_statistical.py

@@ -4,7 +4,7 @@
 from ._lib import _compat
 from ._lib._typing import Array, ArrayNamespace
 
-__all__ = ["cov", "nanmax", "nanmin", "nansum"]
+__all__ = ["cov", "nanmax", "nanmean", "nanmin", "nansum"]
 
 
 def cov(
@@ -321,6 +321,58 @@ def nanmax(
     return _agnostic._statistical.nanmax(a, axis=axis, xp=xp)
 
 
+def nanmean(
+    a: Array,
+    /,
+    *,
+    axis: int | tuple[int, ...] | None = None,
+    xp: ArrayNamespace | None = None,
+) -> Array:
+    """
+    Return the mean of the array elements along a given axis, ignoring NaNs.
+
+    Parameters
+    ----------
+    a : Array
+        Input array.
+    axis : int or tuple of ints or None, optional
+        Axis or axes along which the mean is computed. The default is to compute
+        the mean of the flattened array.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `a`. Default: infer.
+
+    Returns
+    -------
+    array
+        An array of mean values along the given axis, ignoring NaNs.
+
+    Examples
+    --------
+    >>> import array_api_extra as xpx
+    >>> import array_api_strict as xp
+    >>> a = xp.asarray([[5, 3, xp.nan, 1], [4, xp.nan, 2, xp.nan]])
+    >>> xpx.nanmean(a)
+    Array(3., dtype=array_api_strict.float64)
+    >>> xpx.nanmean(a, axis=0)
+    Array([4.5, 3. , 2. , 1. ], dtype=array_api_strict.float64)
+    >>> xpx.nanmean(a, axis=1)
+    Array([3., 3.], dtype=array_api_strict.float64)
+    """
+    if xp is None:
+        xp = _compat.array_namespace(a)
+
+    if (
+        _compat.is_numpy_namespace(xp)
+        or _compat.is_cupy_namespace(xp)
+        or _compat.is_dask_namespace(xp)
+        or _compat.is_jax_namespace(xp)
+        or _compat.is_torch_namespace(xp)
+    ):
+        return xp.nanmean(a, axis=axis)
+
+    return _agnostic._statistical.nanmean(a, axis=axis, xp=xp)
+
+
 def nansum(
     a: Array,
     /,

tests/main/test_statistical.py

@@ -5,14 +5,15 @@
 import numpy as np
 import pytest
 
-from array_api_extra import cov, nanmax, nanmin, nansum
+from array_api_extra import cov, nanmax, nanmean, nanmin, nansum
 from array_api_extra._lib._backends import Backend
 from array_api_extra._lib._compat import array_namespace
 from array_api_extra._lib._compat import device as get_device
 from array_api_extra._lib._typing import Array, ArrayNamespace, Device
 from array_api_extra.testing import assert_close, assert_equal, lazy_xp_function
 
 lazy_xp_function(cov)
+lazy_xp_function(nanmean)
 lazy_xp_function(nansum)
 
 
@@ -506,3 +507,66 @@ def test_xp(self, axis: int | None, expected_list: list[float], xp: ArrayNamespa
         res = nansum(a, axis=axis, xp=xp)
         expected = xp.asarray(expected_list)
         assert_equal(res, expected)
+
+
+class TestNanMean:
+    def test_simple(self, xp: ArrayNamespace):
+        a = xp.asarray([[1.0, 2.0], [3.0, xp.nan]])
+
+        res = nanmean(a)
+        assert res == 2.0
+
+        res = nanmean(a, axis=0)
+        expected = xp.asarray([2.0, 2.0])
+        assert_equal(res, expected)
+
+        res = nanmean(a, axis=1)
+        expected = xp.asarray([1.5, 3.0])
+        assert_equal(res, expected)
+
+    def test_bigger(self, xp: ArrayNamespace):
+        a = xp.asarray(
+            [
+                [1.0, xp.nan, 4.0, 5.0],
+                [xp.nan, -2.0, xp.nan, -4.0],
+                [2.0, 1.0, 3.0, xp.nan],
+            ]
+        )
+
+        res = nanmean(a, axis=0)
+        expected = xp.asarray([1.5, -0.5, 3.5, 0.5])
+        assert_equal(res, expected)
+
+        res = nanmean(a, axis=1)
+        expected = xp.asarray([3.3333333, -3.0, 2.0])
+        assert_close(res, expected)
+
+    @pytest.mark.filterwarnings("ignore:.*Mean of empty slice.*:RuntimeWarning")
+    def test_all_nan_slice(self, xp: ArrayNamespace):
+        a = xp.asarray([[xp.nan, 1.0], [xp.nan, 3.0], [xp.nan, xp.nan]])
+
+        res = nanmean(a, axis=0, xp=xp)
+        expected = xp.asarray([xp.nan, 2.0])
+        assert_equal(res, expected)
+
+    def test_scalar(self, xp: ArrayNamespace):
+        a = xp.asarray(1.0)
+        assert nanmean(a) == 1.0
+
+    @pytest.mark.skip_xp_backend(
+        Backend.TORCH, reason="torch.nanmean does not support tensors on meta device"
+    )
+    @pytest.mark.parametrize("axis", [None, 0, 1])
+    def test_device(self, axis: int | None, xp: ArrayNamespace, device: Device):
+        a = xp.asarray([[4.0, xp.nan, 1.0], [2.0, 5.0, xp.nan]], device=device)
+        res = nanmean(a, axis=axis)
+        assert get_device(res) == device
+
+    @pytest.mark.parametrize(
+        ("axis", "expected_list"), [(0, [3.0, 5.0, 1.0]), (1, [2.5, 3.5])]
+    )
+    def test_xp(self, axis: int | None, expected_list: list[float], xp: ArrayNamespace):
+        a = xp.asarray([[4.0, xp.nan, 1.0], [2.0, 5.0, xp.nan]])
+        res = nanmean(a, axis=axis, xp=xp)
+        expected = xp.asarray(expected_list)
+        assert_equal(res, expected)