ENH: add nanmax

docs/api-assorted.md

@@ -21,6 +21,7 @@
     isin
     kron
     nan_to_num
+    nanmax
     nanmin
     nunique
     one_hot

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",

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)

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

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)