[antonwolfy]

dpnp.nansum, dpnp.nanprod, dpnp.nancumsum and dpnp.nancumprod replace NaN values with dpnp.nan_to_num before reducing, which copies the whole input. For a broadcast view this materializes every repeated element: a (2048, 1, 1024) float32 array broadcast to (2048, 256, 1024) becomes a 2 GiB copy, so the slow TestNansumNanprodHuge tests failed with USMAllocationError on devices whose maximum allocation size is smaller than that.

This PR replaces NaN values only in the distinct elements of a broadcast input, by slicing each zero-stride axis to length 1, and broadcasts the result back to the input shape as a view; the result is only read by the reduction that follows. Inputs without broadcast axes go through dpnp.nan_to_num exactly as before: the only extra work is a check of the shape and strides, which takes 0.5-0.8 µs.

Measured against master on Intel Arc Graphics (Lunar Lake, Level Zero) and on its CPU (OpenCL).

Peak GPU memory, from the per-process drm-total-gtt counter of the xe driver, each case in a fresh process:

Case (float32)	master	this PR
11 cases without broadcast axes (contiguous, strided, transposed, with NaNs, complex64, int32; nansum, nanprod, nancumsum)	same	same (within 0.1 MiB)
nansum/nanprod of (2048, 1→256, 1024), any axis	2055-2063 MiB	14-22 MiB
nansum of (4096, 1→4096), either axis	71 MiB	6 MiB
nancumsum of (512, 1→256, 1024), axis=2	1031 MiB	520 MiB (the output itself is 512 MiB)

Time without broadcast axes, median of 101 interleaved runs with the order alternated, ratio this PR / master:

Case (float32)	GPU	CPU
nansum of 1k / 64k elements	1.03 / 1.05	1.06 / 1.02
nansum / nanprod / nancumsum of 16M elements	0.98 / 0.77 / 1.01	0.98 / 1.01 / 1.00
nansum of a strided 16M view	1.00	1.00

The deviations from 1.00 are within the run-to-run noise; the added check costs under 1 µs.

Time with broadcast axes, median of 31 interleaved runs:

Case (float32)	GPU master → this PR	CPU master → this PR
nansum of (2048, 1→256, 1024), axis=2	296 → 33 ms	1091 → 91 ms
nansum of (2048, 1→256, 1024), axis=1	377 → 78 ms	939 → 118 ms
nansum of (2048, 1→256, 1024), all axes	358 → 101 ms	921 → 409 ms
nanprod of (2048, 1→256, 1024), axis=2	429 → 60 ms	1097 → 92 ms
nansum of (4096, 1→4096), axis=1 / axis=0	12.5 / 10.7 → 1.2 / 1.2 ms	29.7 / 54.7 → 3.6 / 3.9 ms
nancumsum of (512, 1→256, 1024), axis=2	136 → 74 ms	533 → 303 ms

New test_broadcast tests in TestNanProdSum and TestNanCumSumProd check the results on inputs broadcast along different axes. dpnp/tests/test_nanfunctions.py, the third-party sum and prod tests and the slow TestNansumNanprodHuge tests pass on GPU and CPU; the latter failed before on a GPU with a 1 GiB maximum allocation size.

• Have you provided a meaningful PR description?
• Have you added a test, reproducer or referred to an issue with a reproducer?
• Have you tested your changes locally for CPU and GPU devices?
• Have you made sure that new changes do not introduce compiler warnings?
• Have you checked performance impact of proposed changes?
• Have you added documentation for your changes, if necessary?
• Have you added your changes to the changelog?

[github-actions]

View rendered docs @ https://intelpython.github.io/dpnp/index.html

[github-actions]

Array API standard conformance tests for dpnp=0.21.0dev15=np2py314h8d9cdd5_10 ran successfully.
Passed: 1376
Failed: 0
Skipped: 6

[coveralls]

No base build to compare — nansum-no-full-copy into master

[vlad-perevezentsev]

LGTM
Thank you @antonwolfy