Extend nan_to_num with broadcast support of nan, posinf, and neginf (#2754)

The PR extends implementation of `dpnp.nan_to_num` function to align
with NumPy and CuPy which supports `nan`, `posinf`, and `neginf`
keywords as any array through broadcasting.

This PR adds handling for a common path where at least one of the
keywords has non-scalar value.
The path does not assume a dedicated SYCL kernel, instead proposes to
rely on implementation through existing python functions. That can be
improved in the future if required.

Author: antonwolfy

CHANGELOG.md

@@ -48,6 +48,7 @@ Also, that release drops support for Python 3.9, making Python 3.10 the minimum
 * Improved documentation of `file` argument in `dpnp.fromfile` [#2745](https://github.com/IntelPython/dpnp/pull/2745)
 * Aligned `dpnp.trim_zeros` with NumPy 2.4 to support a tuple of integers passed with `axis` keyword [#2746](https://github.com/IntelPython/dpnp/pull/2746)
 * Aligned `strides` property of `dpnp.ndarray` with NumPy and CuPy implementations [#2747](https://github.com/IntelPython/dpnp/pull/2747)
+* Extended `dpnp.nan_to_num` to support broadcasting of `nan`, `posinf`, and `neginf` keywords [#2754](https://github.com/IntelPython/dpnp/pull/2754)
 
 ### Deprecated
 

dpnp/dpnp_iface_mathematical.py

@@ -3646,20 +3646,24 @@ def nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None):
         an array does not require a copy.
 
         Default: ``True``.
-    nan : {int, float, bool}, optional
-        Value to be used to fill ``NaN`` values.
+    nan : {scalar, array_like}, optional
+        Values to be used to fill ``NaN`` values. If no values are passed then
+        ``NaN`` values will be replaced with ``0.0``.
+        Expected to have a real-valued data type for the values.
 
         Default: ``0.0``.
-    posinf : {int, float, bool, None}, optional
-        Value to be used to fill positive infinity values. If no value is
+    posinf : {None, scalar, array_like}, optional
+        Values to be used to fill positive infinity values. If no values are
         passed then positive infinity values will be replaced with a very
         large number.
+        Expected to have a real-valued data type for the values.
 
         Default: ``None``.
-    neginf : {int, float, bool, None} optional
-        Value to be used to fill negative infinity values. If no value is
+    neginf : {None, scalar, array_like}, optional
+        Values to be used to fill negative infinity values. If no values are
         passed then negative infinity values will be replaced with a very
         small (or negative) number.
+        Expected to have a real-valued data type for the values.
 
         Default: ``None``.
 
@@ -3687,13 +3691,22 @@ def nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None):
     array(-1.79769313e+308)
     >>> np.nan_to_num(np.array(np.nan))
     array(0.)
+
     >>> x = np.array([np.inf, -np.inf, np.nan, -128, 128])
     >>> np.nan_to_num(x)
     array([ 1.79769313e+308, -1.79769313e+308,  0.00000000e+000,
            -1.28000000e+002,  1.28000000e+002])
     >>> np.nan_to_num(x, nan=-9999, posinf=33333333, neginf=33333333)
     array([ 3.3333333e+07,  3.3333333e+07, -9.9990000e+03, -1.2800000e+02,
             1.2800000e+02])
+
+    >>> nan = np.array([11, 12, -9999, 13, 14])
+    >>> posinf = np.array([33333333, 11, 12, 13, 14])
+    >>> neginf = np.array([11, 33333333, 12, 13, 14])
+    >>> np.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf)
+    array([ 3.3333333e+07,  3.3333333e+07, -9.9990000e+03, -1.2800000e+02,
+            1.2800000e+02])
+
     >>> y = np.array([complex(np.inf, np.nan), np.nan, complex(np.nan, np.inf)])
     >>> np.nan_to_num(y)
     array([1.79769313e+308 +0.00000000e+000j, # may vary
@@ -3706,33 +3719,32 @@ def nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None):
 
     dpnp.check_supported_arrays_type(x)
 
-    # Python boolean is a subtype of an integer
-    # so additional check for bool is not needed.
-    if not isinstance(nan, (int, float)):
-        raise TypeError(
-            "nan must be a scalar of an integer, float, bool, "
-            f"but got {type(nan)}"
-        )
-    x_type = x.dtype.type
+    def _check_nan_inf(val, val_dt):
+        # Python boolean is a subtype of an integer
+        if not isinstance(val, (int, float)):
+            val = dpnp.asarray(
+                val, dtype=val_dt, sycl_queue=x.sycl_queue, usm_type=x.usm_type
+            )
+        return val
 
-    if not issubclass(x_type, dpnp.inexact):
+    x_type = x.dtype.type
+    if not dpnp.issubdtype(x_type, dpnp.inexact):
         return dpnp.copy(x) if copy else dpnp.get_result_array(x)
 
     max_f, min_f = _get_max_min(x.real.dtype)
+
+    # get dtype of nan and infs values if casting required
+    is_complex = dpnp.issubdtype(x_type, dpnp.complexfloating)
+    if is_complex:
+        val_dt = x.real.dtype
+    else:
+        val_dt = x.dtype
+
+    nan = _check_nan_inf(nan, val_dt)
     if posinf is not None:
-        if not isinstance(posinf, (int, float)):
-            raise TypeError(
-                "posinf must be a scalar of an integer, float, bool, "
-                f"or be None, but got {type(posinf)}"
-            )
-        max_f = posinf
+        max_f = _check_nan_inf(posinf, val_dt)
     if neginf is not None:
-        if not isinstance(neginf, (int, float)):
-            raise TypeError(
-                "neginf must be a scalar of an integer, float, bool, "
-                f"or be None, but got {type(neginf)}"
-            )
-        min_f = neginf
+        min_f = _check_nan_inf(neginf, val_dt)
 
     if copy:
         out = dpnp.empty_like(x)
@@ -3741,19 +3753,45 @@ def nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None):
             raise ValueError("copy is required for read-only array `x`")
         out = x
 
-    x_ary = dpnp.get_usm_ndarray(x)
-    out_ary = dpnp.get_usm_ndarray(out)
+    # handle a special case when nan and infs are all scalars
+    if all(dpnp.isscalar(el) for el in (nan, max_f, min_f)):
+        x_ary = dpnp.get_usm_ndarray(x)
+        out_ary = dpnp.get_usm_ndarray(out)
+
+        q = x.sycl_queue
+        _manager = dpu.SequentialOrderManager[q]
+
+        h_ev, comp_ev = ufi._nan_to_num(
+            x_ary,
+            nan,
+            max_f,
+            min_f,
+            out_ary,
+            q,
+            depends=_manager.submitted_events,
+        )
 
-    q = x.sycl_queue
-    _manager = dpu.SequentialOrderManager[q]
+        _manager.add_event_pair(h_ev, comp_ev)
 
-    h_ev, comp_ev = ufi._nan_to_num(
-        x_ary, nan, max_f, min_f, out_ary, q, depends=_manager.submitted_events
-    )
+        return dpnp.get_result_array(out)
 
-    _manager.add_event_pair(h_ev, comp_ev)
-
-    return dpnp.get_result_array(out)
+    # handle a common case with broadcasting of input nan and infs
+    if is_complex:
+        parts = (x.real, x.imag)
+        parts_out = (out.real, out.imag)
+    else:
+        parts = (x,)
+        parts_out = (out,)
+
+    for part, part_out in zip(parts, parts_out):
+        nan_mask = dpnp.isnan(part)
+        posinf_mask = dpnp.isposinf(part)
+        neginf_mask = dpnp.isneginf(part)
+
+        part = dpnp.where(nan_mask, nan, part, out=part_out)
+        part = dpnp.where(posinf_mask, max_f, part, out=part_out)
+        part = dpnp.where(neginf_mask, min_f, part, out=part_out)
+    return out
 
 
 _NEGATIVE_DOCSTRING = """

dpnp/tests/test_mathematical.py

@@ -1480,37 +1480,35 @@ def test_boolean_array(self):
         expected = numpy.nan_to_num(a)
         assert_allclose(result, expected)
 
-    def test_errors(self):
-        ia = dpnp.array([0, 1, dpnp.nan, dpnp.inf, -dpnp.inf])
-
-        # unsupported type `a`
-        a = dpnp.asnumpy(ia)
-        assert_raises(TypeError, dpnp.nan_to_num, a)
-
-        # unsupported type `nan`
-        i_nan = dpnp.array(1)
-        assert_raises(TypeError, dpnp.nan_to_num, ia, nan=i_nan)
+    @pytest.mark.parametrize("dt", get_float_complex_dtypes())
+    @pytest.mark.parametrize("kw_name", ["nan", "posinf", "neginf"])
+    @pytest.mark.parametrize("val", [[1, 2, -1, -2, 7], (7.0,), numpy.array(1)])
+    def test_nan_infs_array_like(self, dt, kw_name, val):
+        a = numpy.array([0, 1, dpnp.nan, dpnp.inf, -dpnp.inf], dtype=dt)
+        ia = dpnp.array(a)
 
-        # unsupported type `posinf`
-        i_posinf = dpnp.array(1)
-        assert_raises(TypeError, dpnp.nan_to_num, ia, posinf=i_posinf)
+        result = dpnp.nan_to_num(ia, **{kw_name: val})
+        expected = numpy.nan_to_num(a, **{kw_name: val})
+        assert_allclose(result, expected)
 
-        # unsupported type `neginf`
-        i_neginf = dpnp.array(1)
-        assert_raises(TypeError, dpnp.nan_to_num, ia, neginf=i_neginf)
+    @pytest.mark.parametrize("xp", [dpnp, numpy])
+    @pytest.mark.parametrize("kw_name", ["nan", "posinf", "neginf"])
+    def test_nan_infs_complex_dtype(self, xp, kw_name):
+        ia = xp.array([0, 1, xp.nan, xp.inf, -xp.inf])
+        with pytest.raises(TypeError, match="complex"):
+            xp.nan_to_num(ia, **{kw_name: 1j})
 
-    @pytest.mark.parametrize("kwarg", ["nan", "posinf", "neginf"])
-    @pytest.mark.parametrize("value", [1 - 0j, [1, 2], (1,)])
-    def test_errors_diff_types(self, kwarg, value):
-        ia = dpnp.array([0, 1, dpnp.nan, dpnp.inf, -dpnp.inf])
-        with pytest.raises(TypeError):
-            dpnp.nan_to_num(ia, **{kwarg: value})
+    def test_numpy_input_array(self):
+        a = numpy.array([0, 1, dpnp.nan, dpnp.inf, -dpnp.inf])
+        with pytest.raises(TypeError, match="must be any of supported type"):
+            dpnp.nan_to_num(a)
 
-    def test_error_readonly(self):
-        a = dpnp.array([0, 1, dpnp.nan, dpnp.inf, -dpnp.inf])
-        a.flags.writable = False
-        with pytest.raises(ValueError):
-            dpnp.nan_to_num(a, copy=False)
+    @pytest.mark.parametrize("xp", [dpnp, numpy])
+    def test_error_readonly(self, xp):
+        a = xp.array([0, 1, xp.nan, xp.inf, -xp.inf])
+        a.flags["W"] = False
+        with pytest.raises(ValueError, match="read-only"):
+            xp.nan_to_num(a, copy=False)
 
     @pytest.mark.parametrize("copy", [True, False])
     @pytest.mark.parametrize("dt", get_all_dtypes(no_bool=True, no_none=True))
@@ -1522,9 +1520,9 @@ def test_strided(self, copy, dt):
         if dt.kind in "fc":
             a[::4] = numpy.nan
             ia[::4] = dpnp.nan
+
         result = dpnp.nan_to_num(ia[::-2], copy=copy, nan=57.0)
         expected = numpy.nan_to_num(a[::-2], copy=copy, nan=57.0)
-
         assert_dtype_allclose(result, expected)
 
 

dpnp/tests/third_party/cupy/math_tests/test_misc.py

@@ -1,8 +1,10 @@
+from __future__ import annotations
+
 import numpy
 import pytest
 
 import dpnp as cupy
-from dpnp.tests.helper import has_support_aspect64
+from dpnp.tests.helper import has_support_aspect64, numpy_version
 from dpnp.tests.third_party.cupy import testing
 
 
@@ -155,10 +157,7 @@ def test_external_clip4(self, dtype):
             # (min or max) as a keyword argument according to Python Array API.
             # In older versions of numpy, both arguments must be positional;
             # passing only one raises a TypeError.
-            if (
-                xp is numpy
-                and numpy.lib.NumpyVersion(numpy.__version__) < "2.1.0"
-            ):
+            if xp is numpy and numpy_version() < "2.1.0":
                 with pytest.raises(TypeError):
                     xp.clip(a, 3)
             else:
@@ -257,9 +256,10 @@ def test_nan_to_num_inf(self):
     def test_nan_to_num_nan(self):
         self.check_unary_nan("nan_to_num")
 
-    @testing.numpy_cupy_allclose(atol=1e-5, type_check=has_support_aspect64())
+    @pytest.mark.skip("no scalar support")
+    @testing.numpy_cupy_allclose(atol=1e-5)
     def test_nan_to_num_scalar_nan(self, xp):
-        return xp.nan_to_num(xp.array(xp.nan))
+        return xp.nan_to_num(xp.nan)
 
     @pytest.mark.filterwarnings("ignore::RuntimeWarning")
     def test_nan_to_num_inf_nan(self):
@@ -286,14 +286,44 @@ def test_nan_to_num_inplace(self, xp):
         return y
 
     @pytest.mark.parametrize("kwarg", ["nan", "posinf", "neginf"])
-    def test_nan_to_num_broadcast(self, kwarg):
+    @testing.numpy_cupy_array_equal()
+    def test_nan_to_num_broadcast_same_shapes(self, xp, kwarg):
+        x = xp.asarray(
+            [[0, 1, xp.nan, 4], [11, xp.inf, 12, 13]],
+            dtype=cupy.default_float_type(),
+        )
+        y = xp.zeros((2, 4), dtype=x.dtype)
+        return xp.nan_to_num(x, **{kwarg: y})
+
+    @pytest.mark.parametrize("kwarg", ["nan", "posinf", "neginf"])
+    @testing.numpy_cupy_array_equal()
+    def test_nan_to_num_broadcast_different_columns(self, xp, kwarg):
+        x = xp.asarray(
+            [[0, 1, xp.nan, 4], [11, xp.inf, 12, 13]],
+            dtype=cupy.default_float_type(),
+        )
+        y = xp.zeros((2, 1), dtype=x.dtype)
+        return xp.nan_to_num(x, **{kwarg: y})
+
+    @pytest.mark.parametrize("kwarg", ["nan", "posinf", "neginf"])
+    @testing.numpy_cupy_array_equal()
+    def test_nan_to_num_broadcast_different_rows(self, xp, kwarg):
+        x = xp.asarray(
+            [[0, 1, xp.nan, 4], [11, -xp.inf, 12, 13]],
+            dtype=cupy.default_float_type(),
+        )
+        y = xp.zeros((1, 4), dtype=x.dtype)
+        return xp.nan_to_num(x, **{kwarg: y})
+
+    @pytest.mark.parametrize("kwarg", ["nan", "posinf", "neginf"])
+    def test_nan_to_num_broadcast_invalid_shapes(self, kwarg):
         for xp in (numpy, cupy):
             x = xp.asarray([0, 1, xp.nan, 4], dtype=cupy.default_float_type())
-            y = xp.zeros((2, 4), dtype=cupy.default_float_type())
-            with pytest.raises((ValueError, TypeError)):
+            y = xp.zeros((2, 4), dtype=x.dtype)
+            with pytest.raises(ValueError):
                 xp.nan_to_num(x, **{kwarg: y})
-            with pytest.raises((ValueError, TypeError)):
-                xp.nan_to_num(0.0, **{kwarg: y})
+            with pytest.raises(ValueError):
+                xp.nan_to_num(xp.array(0.0), **{kwarg: y})
 
     @testing.for_all_dtypes(no_bool=True, no_complex=True)
     @testing.numpy_cupy_array_equal()