Reducing along a chunked dimension with apply_ufunc and dask="parallelized"

[sadsimulation]

Hello,
I have a function that processes chunks of data on the GPU using cupy by calling a function fun using apply_ufunc.
The input arrays to the function are chunked into dask arrays because my GPU has very limited memory resources.

The function fun could efficiently perform a reduction across one of the chunked dimensions, essentially reducing the size of this dim_0 from N_chunk to 1 (or zero by dropping it entirely).

Usually, apply_ufunc requires dimensions that change size to be core dimensions, which also clashes with the chunking along this dimension. Is there a way to tell apply_ufunc that one of the chunked dimensions will not be present in the output?

[TomNicholas]

Is there a way to tell apply_ufunc that one of the chunked dimensions will not be present in the output?

Can you not just leave that dimension out of output_core_dims?

[sadsimulation]

Hi, I already left it out of output_core_dims, but I still get the error that I would have to allow_rechunk=True. I've done some more digging and this seems to be a limitation of the underlying dask.apply_gufunc.

[dcherian]

The implementation of "parallelized" is only appropriate if you have a single chunk along core dimensions: https://tutorial.xarray.dev/advanced/apply_ufunc/dask_apply_ufunc.html#understanding-dask-parallelized

You'll have to handle the chunkiness manually and use dask="allowed"

[sadsimulation]

Thanks for the reply. I've tried to handle the chunking using the dask primitives directly but for my use case that seemed not to work so well because the chunking limitations seem inherent to dask. Also, creating many small dask chunks just completely kill the performance because the scheduler can't keep up.

[dcherian]

What is this function you're applying