diff --git a/docs/api-assorted.md b/docs/api-assorted.md index 9a69ba23..9a11847e 100644 --- a/docs/api-assorted.md +++ b/docs/api-assorted.md @@ -23,6 +23,7 @@ nan_to_num nanmax nanmin + nansum nunique one_hot pad diff --git a/src/array_api_extra/__init__.py b/src/array_api_extra/__init__.py index 64adcbc8..c7f0e0ad 100644 --- a/src/array_api_extra/__init__.py +++ b/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", diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py index c5c16a37..44d4ccbb 100644 --- a/src/array_api_extra/_delegation.py +++ b/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) diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py index 3f1ef511..ccca2f80 100644 --- a/src/array_api_extra/_lib/_funcs.py +++ b/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) diff --git a/tests/test_funcs.py b/tests/test_funcs.py index 8b639564..cf2a3d26 100644 --- a/tests/test_funcs.py +++ b/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)