diff --git a/docs/api-assorted.md b/docs/api-assorted.md index 7457c4b2..9a69ba23 100644 --- a/docs/api-assorted.md +++ b/docs/api-assorted.md @@ -21,6 +21,7 @@ isin kron nan_to_num + nanmax nanmin nunique one_hot diff --git a/src/array_api_extra/__init__.py b/src/array_api_extra/__init__.py index 578c6373..988321f6 100644 --- a/src/array_api_extra/__init__.py +++ b/src/array_api_extra/__init__.py @@ -13,6 +13,7 @@ isin, kron, nan_to_num, + nanmax, nanmin, one_hot, pad, @@ -55,6 +56,7 @@ "kron", "lazy_apply", "nan_to_num", + "nanmax", "nanmin", "nunique", "one_hot", diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py index 0edcc7cd..f96a9ee0 100644 --- a/src/array_api_extra/_delegation.py +++ b/src/array_api_extra/_delegation.py @@ -35,6 +35,7 @@ "isin", "kron", "nan_to_num", + "nanmax", "nanmin", "one_hot", "pad", @@ -1635,3 +1636,54 @@ def nanmin( return xp.nanmin(a, axis=axis) return _funcs.nanmin(a, axis=axis, xp=xp) + + +def nanmax( + a: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + xp: ModuleType | None = None, +) -> Array: + """ + Return the maximum 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 maximum is computed. The default is to compute + the maximum of the flattened array. + xp : array_namespace, optional + The standard-compatible namespace for `a`. Default: infer. + + Returns + ------- + array + An array of maximum 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, 6], [4, xp.nan, 2, xp.nan]]) + >>> xpx.nanmax(a) + Array(6., dtype=array_api_strict.float64) + >>> xpx.nanmax(a, axis=0) + Array([5., 3., 2., 6.], dtype=array_api_strict.float64) + >>> xpx.nanmax(a, axis=1) + Array([6., 4.], 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) + ): + return xp.nanmax(a, axis=axis) + + return _funcs.nanmax(a, axis=axis, xp=xp) diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py index 7f28cbd2..47333497 100644 --- a/src/array_api_extra/_lib/_funcs.py +++ b/src/array_api_extra/_lib/_funcs.py @@ -38,6 +38,7 @@ "isin", "kron", "nan_to_num", + "nanmax", "nanmin", "nunique", "one_hot", @@ -845,3 +846,24 @@ def nanmin( # 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 nanmax( # numpydoc ignore=PR01,RT01 + a: Array, + /, + *, + axis: int | tuple[int, ...] | None, + xp: ModuleType, +) -> Array: + """See docstring in `array_api_extra._delegation.py`.""" + mask = xp.isnan(a) + device_a = _compat.device(a) + x = xp.max( + xp.where(mask, xp.asarray(-xp.inf, dtype=a.dtype, device=device_a), a), + axis=axis, + ) + # Replace Infs from all NaN slices with NaN again + mask = xp.all(mask, axis=axis) + if xp.any(mask): + x = xp.where(mask, xp.asarray(xp.nan, dtype=x.dtype, device=device_a), x) + return x diff --git a/tests/test_funcs.py b/tests/test_funcs.py index c4e05683..66686d42 100644 --- a/tests/test_funcs.py +++ b/tests/test_funcs.py @@ -30,6 +30,7 @@ isin, kron, nan_to_num, + nanmax, nanmin, nunique, one_hot, @@ -2285,3 +2286,84 @@ def test_xp(self, axis: int | None, expected_list: list[float], xp: ModuleType): res = nanmin(a, axis=axis, xp=xp) expected = xp.asarray(expected_list) assert_equal(res, expected) + + +class TestNanMax: + def test_simple(self, xp: ModuleType): + a = xp.asarray([[5, 3], [6, xp.nan]]) + + # with the default `axis=None` a single scalar is returned + res = nanmax(a) + expected = 6.0 + assert res == expected + + res = nanmax(a, axis=0) + expected = xp.asarray([6.0, 3.0]) + assert_equal(res, expected) + + res = nanmax(a, axis=1) + expected = xp.asarray([5.0, 6.0]) + assert_equal(res, expected) + + def test_bigger(self, xp: ModuleType): + a = xp.asarray( + [ + [1, xp.nan, 4, 5], + [xp.nan, 2, xp.nan, 4], + [6, 1, 3, xp.nan], + ] + ) + + res = nanmax(a, axis=0) + expected = xp.asarray([6.0, 2.0, 4.0, 5.0]) + assert_equal(res, expected) + + res = nanmax(a, axis=1) + expected = xp.asarray([5.0, 4.0, 6.0]) + assert_equal(res, expected) + + def test_with_infinity(self, xp: ModuleType): + a = xp.asarray([0.1, 5.0, xp.nan, -xp.inf]) + res = nanmax(a) + expected = 5.0 + assert res == expected + + a = xp.asarray([3.0, 10.0, xp.nan, xp.inf]) + res = nanmax(a) + expected = xp.inf + assert res == expected + + def test_scalar(self, xp: ModuleType): + a = xp.asarray(1.0) + assert nanmax(a) == 1.0 + + @pytest.mark.filterwarnings("ignore:.*All-NaN slice*.:RuntimeWarning") + def test_all_nan_slice_2d(self, xp: ModuleType): + a = xp.asarray( + [ + [xp.nan, 5.0], + [xp.nan, 2.0], + ] + ) + + res = nanmax(a, axis=0, xp=xp) + expected = xp.asarray([xp.nan, 5.0]) + assert_equal(res, expected) + + @pytest.mark.skip_xp_backend( + Backend.TORCH, reason="torch.nanmax does not support tensors on meta device" + ) + @pytest.mark.parametrize("axis", [None, 0, 1]) + def test_device(self, axis: int | None, xp: ModuleType, device: Device): + a = xp.asarray([[4, xp.nan, 1], [2, 5, xp.nan]], device=device) + res = nanmax(a, axis=axis) + assert get_device(res) == device + + @pytest.mark.parametrize( + ("axis", "expected_list"), [(0, [4.0, 3.0, 1.0]), (1, [4.0, 3.0])] + ) + def test_xp(self, axis: int | None, expected_list: list[float], xp: ModuleType): + a = xp.asarray([[4, xp.nan, 1], [2, 3, xp.nan]]) + res = nanmax(a, axis=axis, xp=xp) + expected = xp.asarray(expected_list) + assert_equal(res, expected)