Avoid a full-size copy of broadcast input in NaN reductions

CHANGELOG.md

@@ -58,6 +58,7 @@ This release is compatible with NumPy 2.5.
 * Changed `dpnp.sort`, `dpnp.argsort`, and their `dpnp.ndarray`/`dpnp.tensor` counterparts to place `NaN` values last instead of first when sorting in descending order [#3066](https://github.com/IntelPython/dpnp/pull/3066)
 * Documented and made explicit support for `copy=None` ("copy only if needed") in `dpnp.nan_to_num` [#3083](https://github.com/IntelPython/dpnp/pull/3083)
 * Changed `dpnp.linalg.pinv` to compute the default tolerance from the machine epsilon of the data type used for the computation, and to return a floating-point result for an empty integer or boolean array [#3090](https://github.com/IntelPython/dpnp/pull/3090)
+* Avoided a full-size copy of a broadcast input in `dpnp.nansum`, `dpnp.nanprod`, `dpnp.nancumsum` and `dpnp.nancumprod` by replacing `NaN` values only in the distinct elements of the input [#3108](https://github.com/IntelPython/dpnp/pull/3108)
 
 ### Deprecated
 

dpnp/dpnp_iface_nanfunctions.py

@@ -53,6 +53,10 @@ def _replace_nan_no_mask(a, val):
     the `val` value, and return the copy. If `a` is not of inexact type,
     do nothing and return `a`.
 
+    If `a` is a broadcast view, only its distinct elements are copied and
+    the copy is broadcast back to the shape of `a`, so the result must only
+    be read.
+
     Parameters
     ----------
     a : {dpnp.ndarray, usm_ndarray}
@@ -69,10 +73,17 @@ def _replace_nan_no_mask(a, val):
     """
 
     dpnp.check_supported_arrays_type(a)
-    if dpnp.issubdtype(a.dtype, dpnp.inexact):
+    if not dpnp.issubdtype(a.dtype, dpnp.inexact):
+        return a
+
+    repeated = [st == 0 and sh > 1 for sh, st in zip(a.shape, a.strides)]
+    if not any(repeated):
         return dpnp.nan_to_num(a, nan=val, posinf=dpnp.inf, neginf=-dpnp.inf)
 
-    return a
+    # a zero-stride axis repeats the same elements, keep one of them
+    base = a[tuple(slice(0, 1) if r else slice(None) for r in repeated)]
+    res = dpnp.nan_to_num(base, nan=val, posinf=dpnp.inf, neginf=-dpnp.inf)
+    return dpnp.broadcast_to(res, a.shape)
 
 
 def _replace_nan(a, val):

dpnp/tests/test_nanfunctions.py

@@ -114,6 +114,23 @@ def test_allnans(self, func, dtype, array):
         # Array API indicates that the behavior is unspecified
         assert_dtype_allclose(result, expected, check_shape=False)
 
+    @pytest.mark.parametrize("dtype", get_float_complex_dtypes())
+    @pytest.mark.parametrize("axis", [None, 0, 1, 2])
+    @pytest.mark.parametrize(
+        "shape", [(1, 3, 4), (2, 1, 4), (2, 3, 1), (1, 3, 1)]
+    )
+    def test_broadcast(self, func, shape, axis, dtype):
+        # positive values avoid cancellation in the cumulative sums
+        a = generate_random_numpy_array(shape, dtype=dtype, low=1)
+        a.flat[::3] = numpy.nan
+        ia = dpnp.array(a)
+
+        a = numpy.broadcast_to(a, (2, 3, 4))
+        ia = dpnp.broadcast_to(ia, (2, 3, 4))
+        result = getattr(dpnp, func)(ia, axis=axis)
+        expected = getattr(numpy, func)(a, axis=axis)
+        assert_dtype_allclose(result, expected)
+
     @pytest.mark.parametrize("axis", [None, 0, 1])
     def test_empty(self, func, axis):
         a = numpy.zeros((0, 3))
@@ -540,6 +557,23 @@ def test_basic(self, func, axis, keepdims, dtype):
 
         assert_allclose(result, expected, rtol=1e-5)
 
+    @pytest.mark.parametrize("dtype", get_float_complex_dtypes())
+    @pytest.mark.parametrize("axis", [None, 0, 1, 2, (0, 2)])
+    @pytest.mark.parametrize(
+        "shape", [(1, 3, 4), (2, 1, 4), (2, 3, 1), (1, 3, 1)]
+    )
+    def test_broadcast(self, func, shape, axis, dtype):
+        # positive values avoid cancellation in the sums
+        a = generate_random_numpy_array(shape, dtype=dtype, low=1)
+        a.flat[::3] = numpy.nan
+        ia = dpnp.array(a)
+
+        a = numpy.broadcast_to(a, (2, 3, 4))
+        ia = dpnp.broadcast_to(ia, (2, 3, 4))
+        result = getattr(dpnp, func)(ia, axis=axis)
+        expected = getattr(numpy, func)(a, axis=axis)
+        assert_dtype_allclose(result, expected)
+
     @pytest.mark.usefixtures("suppress_complex_warning")
     @pytest.mark.usefixtures("suppress_invalid_numpy_warnings")
     @pytest.mark.parametrize("in_dt", get_float_complex_dtypes())