ENH: add nansum

docs/api-assorted.md

@@ -23,6 +23,7 @@
     nan_to_num
     nanmax
     nanmin
+    nansum
     nunique
     one_hot
     pad

src/array_api_extra/__init__.py

@@ -15,6 +15,7 @@
     nan_to_num,
     nanmax,
     nanmin,
+    nansum,
     nunique,
     one_hot,
     pad,
@@ -58,6 +59,7 @@
     "nan_to_num",
     "nanmax",
     "nanmin",
+    "nansum",
     "nunique",
     "one_hot",
     "pad",

src/array_api_extra/_delegation.py

@@ -39,6 +39,7 @@
     "nan_to_num",
     "nanmax",
     "nanmin",
+    "nansum",
     "nunique",
     "one_hot",
     "pad",
@@ -1732,3 +1733,55 @@ def nanmax(
         return xp.nanmax(a, axis=axis)
 
     return _funcs.nanmax(a, axis=axis, xp=xp)
+
+
+def nansum(
+    a: Array,
+    /,
+    *,
+    axis: int | tuple[int, ...] | None = None,
+    xp: ArrayNamespace | None = None,
+) -> Array:
+    """
+    Return the sum 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 sum is computed. The default is to compute
+        the sum of the flattened array.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `a`. Default: infer.
+
+    Returns
+    -------
+    array
+        An array of sum 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.nansum(a)
+    Array(15., dtype=array_api_strict.float64)
+    >>> xpx.nansum(a, axis=0)
+    Array([9., 3., 2., 1.], dtype=array_api_strict.float64)
+    >>> xpx.nansum(a, axis=1)
+    Array([9., 6.], dtype=array_api_strict.float64)
+    """
+    if xp is None:
+        xp = array_namespace(a)
+
+    if (
+        is_numpy_namespace(xp)
+        or is_cupy_namespace(xp)
+        or is_dask_namespace(xp)
+        or is_jax_namespace(xp)
+        or is_torch_namespace(xp)
+    ):
+        return xp.nansum(a, axis=axis)
+
+    return _funcs.nansum(a, axis=axis, xp=xp)

src/array_api_extra/_lib/_funcs.py

@@ -39,6 +39,7 @@
     "nan_to_num",
     "nanmax",
     "nanmin",
+    "nansum",
     "nunique",
     "one_hot",
     "pad",
@@ -842,3 +843,17 @@ def nanmax(  # numpydoc ignore=PR01,RT01
     if xp.any(mask):
         x = xp.where(mask, xp.asarray(xp.nan, dtype=x.dtype, device=device_a), x)
     return x
+
+
+def nansum(  # numpydoc ignore=PR01,RT01
+    a: Array,
+    /,
+    *,
+    axis: int | tuple[int, ...] | None,
+    xp: ArrayNamespace,
+) -> Array:
+    """See docstring in `array_api_extra._delegation.py`."""
+    mask = xp.isnan(a)
+    device_a = _compat.device(a)
+    zero = xp.asarray(0, dtype=a.dtype, device=device_a)
+    return xp.sum(xp.where(mask, zero, a), axis=axis)

tests/test_funcs.py

@@ -31,6 +31,7 @@
     nan_to_num,
     nanmax,
     nanmin,
+    nansum,
     nunique,
     one_hot,
     pad,
@@ -70,6 +71,7 @@
 lazy_xp_function(isin)
 lazy_xp_function(kron)
 lazy_xp_function(nan_to_num)
+lazy_xp_function(nansum)
 lazy_xp_function(nunique)
 lazy_xp_function(one_hot)
 lazy_xp_function(pad)
@@ -2392,3 +2394,66 @@ def test_xp(self, axis: int | None, expected_list: list[float], xp: ArrayNamespa
         res = nanmax(a, axis=axis, xp=xp)
         expected = xp.asarray(expected_list)
         assert_equal(res, expected)
+
+
+class TestNanSum:
+    def test_simple(self, xp: ArrayNamespace):
+        a = xp.asarray([[1.0, 2.0], [3.0, xp.nan]])
+
+        res = nansum(a)
+        expected = 6.0
+        assert res == expected
+
+        res = nansum(a, axis=0)
+        expected = xp.asarray([4.0, 2.0])
+        assert_equal(res, expected)
+
+        res = nansum(a, axis=1)
+        expected = xp.asarray([3.0, 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 = nansum(a, axis=0)
+        expected = xp.asarray([3.0, -1.0, 7.0, 1.0])
+        assert_equal(res, expected)
+
+        res = nansum(a, axis=1)
+        expected = xp.asarray([10.0, -6.0, 6.0])
+        assert_equal(res, expected)
+
+    def test_all_nan_slice(self, xp: ArrayNamespace):
+        a = xp.asarray([[xp.nan, 1.0], [xp.nan, xp.nan]])
+
+        res = nansum(a, axis=0)
+        expected = xp.asarray([0.0, 1.0])
+        assert_equal(res, expected)
+
+    def test_scalar(self, xp: ArrayNamespace):
+        a = xp.asarray(1.0)
+        assert nansum(a) == 1.0
+
+    @pytest.mark.skip_xp_backend(
+        Backend.TORCH, reason="torch.nansum 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 = nansum(a, axis=axis)
+        assert get_device(res) == device
+
+    @pytest.mark.parametrize(
+        ("axis", "expected_list"), [(0, [6.0, 3.0, 1.0]), (1, [5.0, 5.0])]
+    )
+    def test_xp(self, axis: int | None, expected_list: list[float], xp: ArrayNamespace):
+        a = xp.asarray([[4.0, xp.nan, 1.0], [2.0, 3.0, xp.nan]])
+        res = nansum(a, axis=axis, xp=xp)
+        expected = xp.asarray(expected_list)
+        assert_equal(res, expected)