diff --git a/CHANGELOG.md b/CHANGELOG.md index 3d6b8161ef8..9b1c81bd25b 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/dpnp/dpnp_iface_nanfunctions.py b/dpnp/dpnp_iface_nanfunctions.py index b528316b485..f634064b016 100644 --- a/dpnp/dpnp_iface_nanfunctions.py +++ b/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): diff --git a/dpnp/tests/test_nanfunctions.py b/dpnp/tests/test_nanfunctions.py index 810ac9b61af..8f20a294374 100644 --- a/dpnp/tests/test_nanfunctions.py +++ b/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())