implemented nanmean

Author: UshnahAbbasi

docs/api-assorted.md

@@ -0,0 +1,41 @@
+# Assorted functions
+
+```{eval-rst}
+.. currentmodule:: array_api_extra
+.. autosummary::
+    :nosignatures:
+    :toctree: generated
+
+    angle
+    apply_where
+    argpartition
+    at
+    atleast_nd
+    broadcast_shapes
+    cov
+    create_diagonal
+    default_dtype
+    deg2rad
+    diag_indices
+    expand_dims
+    isclose
+    isin
+    kron
+    nan_to_num
+    nanmax
+    nanmean
+    nanmin
+    nansum
+    nunique
+    one_hot
+    pad
+    partition
+    rad2deg
+    searchsorted
+    setdiff1d
+    sinc
+    tril_indices
+    triu_indices
+    union1d
+    unravel_index
+```

src/array_api_extra/__init__.py

@@ -1,19 +1,43 @@
 """Extra array functions built on top of the array API standard."""
 
 from . import testing
-from ._agnostic._elementwise import angle, apply_where
-from ._agnostic._inspection import default_dtype
-from ._at import at
-from ._creation import create_diagonal, one_hot
-from ._elementwise import deg2rad, isclose, nan_to_num, rad2deg, sinc
-from ._indexing import diag_indices, tril_indices, triu_indices, unravel_index
-from ._lazy import lazy_apply
-from ._linalg import kron
-from ._manipulation import atleast_nd, broadcast_shapes, expand_dims, pad
-from ._searching import searchsorted
-from ._set import isin, nunique, setdiff1d, union1d
-from ._sorting import argpartition, partition
-from ._statistical import cov, nanmax, nanmin, nansum
+from ._delegation import (
+    argpartition,
+    atleast_nd,
+    broadcast_shapes,
+    cov,
+    create_diagonal,
+    deg2rad,
+    diag_indices,
+    expand_dims,
+    isclose,
+    isin,
+    kron,
+    nan_to_num,
+    nanmax,
+    nanmean,
+    nanmin,
+    nansum,
+    nunique,
+    one_hot,
+    pad,
+    partition,
+    rad2deg,
+    searchsorted,
+    setdiff1d,
+    sinc,
+    tril_indices,
+    triu_indices,
+    union1d,
+    unravel_index,
+)
+from ._lib._at import at
+from ._lib._funcs import (
+    angle,
+    apply_where,
+    default_dtype,
+)
+from ._lib._lazy import lazy_apply
 
 __version__ = "0.11.2.dev0"
 
@@ -37,6 +61,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,29 @@ 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.count_nonzero(~mask, axis=axis)
+    safe_count = xp.astype(
+        xp.where(count == 0, xp.ones_like(count), count),
+        sum_.dtype,
+        copy=False,
+    )
+    result = sum_ / safe_count
+    if xp.any(count == 0):
+        result = xp.where(
+            count == 0,
+            xp.full_like(result, xp.nan),
+            result,
+        )
+    return result