Implement dpnp.nan_to_num()

dpnp/dpnp_iface_mathematical.py

@@ -2264,12 +2264,12 @@ def modf(x1, **kwargs):
 
 def nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None):
     """
-    Replace NaN with zero and infinity with large finite numbers (default
+    Replace ``NaN`` with zero and infinity with large finite numbers (default
     behaviour) or with the numbers defined by the user using the `nan`,
     `posinf` and/or `neginf` keywords.
 
-    If `x` is inexact, NaN is replaced by zero or by the user defined value in
-    `nan` keyword, infinity is replaced by the largest finite floating point
+    If `x` is inexact, ``NaN`` is replaced by zero or by the user defined value
+    in `nan` keyword, infinity is replaced by the largest finite floating point
     values representable by ``x.dtype`` or by the user defined value in
     `posinf` keyword and -infinity is replaced by the most negative finite
     floating point values representable by ``x.dtype`` or by the user defined
@@ -2287,12 +2287,11 @@ def nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None):
     x : {dpnp.ndarray, usm_ndarray}
         Input data.
     copy : bool, optional
-        Whether to create a copy of `x` (True) or to replace values
-        in-place (False). The in-place operation only occurs if
-        casting to an array does not require a copy.
-        Default: ``True``.
+        Whether to create a copy of `x` (``True``) or to replace values
+        in-place (``False``). The in-place operation only occurs if casting to
+        an array does not require a copy.
     nan : {int, float}, optional
-        Value to be used to fill NaN values.
+        Value to be used to fill ``NaN`` values.
         Default: ``0.0``.
     posinf : {int, float, None}, optional
         Value to be used to fill positive infinity values. If no value is
@@ -2308,8 +2307,8 @@ def nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None):
     Returns
     -------
     out : dpnp.ndarray
-        `x`, with the non-finite values replaced. If `copy` is False, this may
-        be `x` itself.
+        `x`, with the non-finite values replaced. If `copy` is ``False``, this
+        may be `x` itself.
 
     See Also
     --------
@@ -2348,7 +2347,10 @@ def nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None):
 
     dpnp.check_supported_arrays_type(x)
 
-    x = dpnp.array(x, copy=copy)
+    if not dpnp.isscalar(nan):
+        raise TypeError(f"nan must be a scalar, but got {type(nan)}")
+
+    out = dpnp.empty_like(x) if copy else x
     x_type = x.dtype.type
 
     if not issubclass(x_type, dpnp.inexact):
@@ -2357,22 +2359,35 @@ def nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None):
     parts = (
         (x.real, x.imag) if issubclass(x_type, dpnp.complexfloating) else (x,)
     )
+    parts_out = (
+        (out.real, out.imag)
+        if issubclass(x_type, dpnp.complexfloating)
+        else (out,)
+    )
     max_f, min_f = _get_max_min(x.real.dtype)
     if posinf is not None:
+        if not dpnp.isscalar(posinf):
+            raise TypeError(
+                f"posinf must be a scalar or None, but got {type(posinf)}"
+            )
         max_f = posinf
     if neginf is not None:
+        if not dpnp.isscalar(neginf):
+            raise TypeError(
+                f"neginf must be a scalar or None, but got {type(neginf)}"
+            )
         min_f = neginf
 
-    for part in parts:
+    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)
-        part = dpnp.where(posinf_mask, max_f, part, out=part)
-        part = dpnp.where(neginf_mask, min_f, part, out=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 x
+    return out
 
 
 _NEGATIVE_DOCSTRING = """

tests/test_mathematical.py

@@ -1116,6 +1116,59 @@ def test_subtract(self, dtype, lhs, rhs):
         self._test_mathematical("subtract", dtype, lhs, rhs, check_type=False)
 
 
+class TestNanToNum:
+    @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
+    @pytest.mark.parametrize("shape", [(3,), (2, 3), (3, 2, 2)])
+    def test_nan_to_num(self, dtype, shape):
+        a = numpy.random.randn(*shape).astype(dtype)
+        if not dpnp.issubdtype(dtype, dpnp.integer):
+            a.flat[1] = numpy.nan
+        a_dp = dpnp.array(a)
+
+        result = dpnp.nan_to_num(a_dp)
+        expected = numpy.nan_to_num(a)
+        assert_allclose(result, expected)
+
+    @pytest.mark.parametrize(
+        "data", [[], [numpy.nan], [numpy.inf], [-numpy.inf]]
+    )
+    @pytest.mark.parametrize("dtype", get_float_complex_dtypes())
+    def test_empty_and_single_value_arrays(self, data, dtype):
+        a = numpy.array(data, dtype)
+        ia = dpnp.array(a)
+
+        result = dpnp.nan_to_num(ia)
+        expected = numpy.nan_to_num(a)
+        assert_allclose(result, expected)
+
+    def test_boolean_array(self):
+        a = numpy.array([True, False, numpy.nan], dtype=bool)
+        ia = dpnp.array(a)
+
+        result = dpnp.nan_to_num(ia)
+        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_np = dpnp.asnumpy(ia)
+        assert_raises(TypeError, dpnp.nan_to_num, a_np)
+
+        # unsupported type `nan`
+        i_nan = dpnp.array(1)
+        assert_raises(TypeError, dpnp.nan_to_num, ia, nan=i_nan)
+
+        # unsupported type `posinf`
+        i_posinf = dpnp.array(1)
+        assert_raises(TypeError, dpnp.nan_to_num, ia, posinf=i_posinf)
+
+        # unsupported type `neginf`
+        i_neginf = dpnp.array(1)
+        assert_raises(TypeError, dpnp.nan_to_num, ia, neginf=i_neginf)
+
+
 class TestNextafter:
     @pytest.mark.parametrize("dt", get_float_dtypes())
     @pytest.mark.parametrize(
@@ -1445,40 +1498,6 @@ def test_power_scalar(shape, dtype):
     assert_allclose(result, expected, rtol=1e-6)
 
 
-class TestNanToNum:
-    @pytest.mark.parametrize("dtype", get_all_dtypes())
-    @pytest.mark.parametrize("shape", [(3,), (2, 3), (3, 2, 2)])
-    def test_nan_to_num(self, dtype, shape):
-        a = numpy.random.randn(*shape).astype(dtype)
-        if not dpnp.issubdtype(dtype, dpnp.integer):
-            a.flat[1] = numpy.nan
-        a_dp = dpnp.array(a)
-
-        result = dpnp.nan_to_num(a_dp)
-        expected = numpy.nan_to_num(a)
-        assert_allclose(result, expected)
-
-    @pytest.mark.parametrize(
-        "data", [[], [numpy.nan], [numpy.inf], [-numpy.inf]]
-    )
-    @pytest.mark.parametrize("dtype", get_float_complex_dtypes())
-    def test_empty_and_single_value_arrays(self, data, dtype):
-        a = numpy.array(data, dtype)
-        ia = dpnp.array(a)
-
-        result = dpnp.nan_to_num(ia)
-        expected = numpy.nan_to_num(a)
-        assert_allclose(result, expected)
-
-    def test_boolean_array(self):
-        a = numpy.array([True, False, numpy.nan], dtype=bool)
-        ia = dpnp.array(a)
-
-        result = dpnp.nan_to_num(ia)
-        expected = numpy.nan_to_num(a)
-        assert_allclose(result, expected)
-
-
 @pytest.mark.parametrize(
     "data",
     [[[1.0, -1.0], [0.1, -0.1]], [-2, -1, 0, 1, 2]],

tests/test_sycl_queue.py

@@ -2329,3 +2329,20 @@ def test_astype(device_x, device_y):
     sycl_queue = dpctl.SyclQueue(device_y)
     y = dpnp.astype(x, dtype="f4", device=sycl_queue)
     assert_sycl_queue_equal(y.sycl_queue, sycl_queue)
+
+
+@pytest.mark.parametrize("copy", [True, False], ids=["True", "False"])
+@pytest.mark.parametrize(
+    "device",
+    valid_devices,
+    ids=[device.filter_string for device in valid_devices],
+)
+def test_nan_to_num(copy, device):
+    a = dpnp.array([-dpnp.nan, -1, 0, 1, dpnp.nan], device=device)
+    result = dpnp.nan_to_num(a, copy=copy)
+
+    assert_sycl_queue_equal(result.sycl_queue, a.sycl_queue)
+    if copy:
+        assert result is not a
+    else:
+        assert result is a

tests/test_usm_type.py

@@ -1354,3 +1354,16 @@ def test_histogram_bin_edges(usm_type_v, usm_type_w):
     assert v.usm_type == usm_type_v
     assert w.usm_type == usm_type_w
     assert edges.usm_type == du.get_coerced_usm_type([usm_type_v, usm_type_w])
+
+
+@pytest.mark.parametrize("copy", [True, False], ids=["True", "False"])
+@pytest.mark.parametrize("usm_type_a", list_of_usm_types, ids=list_of_usm_types)
+def test_nan_to_num(copy, usm_type_a):
+    a = dp.array([-dp.nan, -1, 0, 1, dp.nan], usm_type=usm_type_a)
+    result = dp.nan_to_num(a, copy=copy)
+
+    assert result.usm_type == usm_type_a
+    if copy:
+        assert result is not a
+    else:
+        assert result is a

tests/third_party/cupy/math_tests/test_misc.py

@@ -273,16 +273,17 @@ def test_nan_to_num_inplace(self, xp):
         assert x is y
         return y
 
+    @pytest.mark.skip(reason="nan, posinf, neginf as array are not supported")
     @pytest.mark.parametrize("kwarg", ["nan", "posinf", "neginf"])
     def test_nan_to_num_broadcast(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):
+            with pytest.raises(TypeError):
                 xp.nan_to_num(x, **{kwarg: y})
             # dpnp.nan_to_num() doesn`t support a scalar as an input
             # convert 0.0 to 0-ndim array
-            with pytest.raises(ValueError):
+            with pytest.raises(TypeError):
                 x_ndim_0 = xp.array(0.0)
                 xp.nan_to_num(x_ndim_0, **{kwarg: y})