Add more testing when input array has any floating or complex dtype

Author: antonwolfy

dpnp/dpnp_iface_mathematical.py

@@ -3719,14 +3719,12 @@ def nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None):
 
     dpnp.check_supported_arrays_type(x)
 
-    def _check_nan_inf(val, name):
+    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, sycl_queue=x.sycl_queue, usm_type=x.usm_type
+                val, dtype=val_dt, sycl_queue=x.sycl_queue, usm_type=x.usm_type
             )
-            if dpnp.issubdtype(val.dtype, dpnp.complexfloating):
-                raise TypeError(f"{name} must not be of a complex type")
         return val
 
     x_type = x.dtype.type
@@ -3735,11 +3733,18 @@ def _check_nan_inf(val, name):
 
     max_f, min_f = _get_max_min(x.real.dtype)
 
-    nan = _check_nan_inf(nan, "nan")
+    # 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:
-        max_f = _check_nan_inf(posinf, "posinf")
+        max_f = _check_nan_inf(posinf, val_dt)
     if neginf is not None:
-        min_f = _check_nan_inf(neginf, "neginf")
+        min_f = _check_nan_inf(neginf, val_dt)
 
     if copy:
         out = dpnp.empty_like(x)
@@ -3771,7 +3776,7 @@ def _check_nan_inf(val, name):
         return dpnp.get_result_array(out)
 
     # handle a common case with broadcasting of input nan and infs
-    if dpnp.issubdtype(x_type, dpnp.complexfloating):
+    if is_complex:
         parts = (x.real, x.imag)
         parts_out = (out.real, out.imag)
     else:

dpnp/tests/test_mathematical.py

@@ -1480,10 +1480,11 @@ def test_boolean_array(self):
         expected = numpy.nan_to_num(a)
         assert_allclose(result, expected)
 
+    @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,), numpy.array(1)])
-    def test_nan_infs_array_like(self, kw_name, val):
-        a = numpy.array([0, 1, dpnp.nan, dpnp.inf, -dpnp.inf])
+    @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)
 
         result = dpnp.nan_to_num(ia, **{kw_name: val})
@@ -1494,7 +1495,7 @@ def test_nan_infs_array_like(self, kw_name, val):
     @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, ValueError), match="complex.*type"):
+        with pytest.raises(TypeError, match="complex"):
             xp.nan_to_num(ia, **{kw_name: 1j})
 
     def test_numpy_input_array(self):