Reuse nan_to_num in nanargmax and nanargmin

`_replace_nan_no_mask` becomes `_replace_nan` and the responsibility of raising the error moves into `_replace_nan_test_axis`

Author: ndgrigorian

dpnp/dpnp_iface_nanfunctions.py

@@ -60,7 +60,7 @@
 ]
 
 
-def _replace_nan_no_mask(a, val):
+def _replace_nan(a, val):
     """
     Replace NaNs in array `a` with `val`.
 
@@ -90,14 +90,14 @@ def _replace_nan_no_mask(a, val):
     return a
 
 
-def _replace_nan(a, val):
+def _replace_nan_test_axis(a, val, axis):
     """
-    Replace NaNs in array `a` with `val`.
+    Replace NaNs in array `a` with `val` and test for all-NaN slices.
 
-    If `a` is of inexact type, make a copy of `a`, replace NaNs with
-    the `val` value, and return the copy together with a boolean mask
-    marking the locations where NaNs were present. If `a` is not of
-    inexact type, do nothing and return `a` together with a mask of None.
+    If `a` is of inexact type, test along an axis for an all-NaN slice.
+    If none are found, make a copy of `a`, replace NaNs with the `val`
+    value, and return the copy. If `a` is not of inexact type, do
+    nothing and return `a`.
 
     Parameters
     ----------
@@ -108,27 +108,21 @@ def _replace_nan(a, val):
 
     Returns
     -------
-    out : {dpnp.ndarray}
+    out : dpnp.ndarray
         If `a` is of inexact type, return a copy of `a` with the NaNs
         replaced by the fill value, otherwise return `a`.
-    mask: {bool, None}
-        If `a` is of inexact type, return a boolean mask marking locations of
-        NaNs, otherwise return ``None``.
 
     """
 
     dpnp.check_supported_arrays_type(a)
     if dpnp.issubdtype(a.dtype, dpnp.inexact):
         mask = dpnp.isnan(a)
-        if not dpnp.any(mask):
-            mask = None
-        else:
-            a = dpnp.array(a, copy=True)
-            dpnp.copyto(a, val, where=mask)
-    else:
-        mask = None
+        mask = dpnp.all(mask, axis=axis)
+        if dpnp.any(mask):
+            raise ValueError("All-NaN slice encountered")
+        return dpnp.nan_to_num(a, nan=val, posinf=dpnp.inf, neginf=-dpnp.inf)
 
-    return a, mask
+    return a
 
 
 def nanargmax(a, axis=None, out=None, *, keepdims=False):
@@ -197,11 +191,7 @@ def nanargmax(a, axis=None, out=None, *, keepdims=False):
 
     """
 
-    a, mask = _replace_nan(a, -dpnp.inf)
-    if mask is not None:
-        mask = dpnp.all(mask, axis=axis)
-        if dpnp.any(mask):
-            raise ValueError("All-NaN slice encountered")
+    a = _replace_nan_test_axis(a, -dpnp.inf, axis)
     return dpnp.argmax(a, axis=axis, out=out, keepdims=keepdims)
 
 
@@ -271,11 +261,7 @@ def nanargmin(a, axis=None, out=None, *, keepdims=False):
 
     """
 
-    a, mask = _replace_nan(a, dpnp.inf)
-    if mask is not None:
-        mask = dpnp.all(mask, axis=axis)
-        if dpnp.any(mask):
-            raise ValueError("All-NaN slice encountered")
+    a = _replace_nan_test_axis(a, dpnp.inf, axis)
     return dpnp.argmin(a, axis=axis, out=out, keepdims=keepdims)
 
 
@@ -345,7 +331,7 @@ def nancumprod(a, axis=None, dtype=None, out=None):
 
     """
 
-    a = _replace_nan_no_mask(a, 1.0)
+    a = _replace_nan(a, 1.0)
     return dpnp.cumprod(a, axis=axis, dtype=dtype, out=out)
 
 
@@ -415,7 +401,7 @@ def nancumsum(a, axis=None, dtype=None, out=None):
 
     """
 
-    a = _replace_nan_no_mask(a, 0.0)
+    a = _replace_nan(a, 0.0)
     return dpnp.cumsum(a, axis=axis, dtype=dtype, out=out)
 
 
@@ -914,7 +900,7 @@ def nanprod(
 
     """
 
-    a = _replace_nan_no_mask(a, 1.0)
+    a = _replace_nan(a, 1.0)
     return dpnp.prod(
         a,
         axis=axis,
@@ -1018,7 +1004,7 @@ def nansum(
 
     """
 
-    a = _replace_nan_no_mask(a, 0.0)
+    a = _replace_nan(a, 0.0)
     return dpnp.sum(
         a,
         axis=axis,