issue_comments

Modify apply_ufunc: remove the check that the input_core_dims must not be chunked add parameter output_chunks

My main concern is ensuring that someone does not inadvertently apply a function not designed for multiple chunks to dask arrays. For example, suppose the function being applied is np.median.

Some loud flag that makes it very obvious what's going on seems like a good idea, e.g., possibly_chunked_core_dims=['x']?

Then we also need some sort of guarantee that chunked core dimensions aren't entirely removed, or else xarray/dask won't know how to stack them back up. I guess we could check to make sure that at least as many output core dimensions appear as appear in inputs cor edimensions?

Try: ```python import dask.array import numpy as np

def mulsum_chunk(a, b): return np.einsum('...i,...i', a, b)[..., np.newaxis]

def mulsum(a, b): # needs broadcasting/rechunking for a,b mapped = dask.array.map_blocks(mulsum_chunk, a, b, dtype=float, chunks=a.chunks[:-1] + (tuple(1 for _ in a.chunks[-1]),)) return dask.array.sum(mapped, axis=-1) ```

could you make an example? That was my first thought but I could not figure out how to make the apply_ufunc do it.

OK, thinking a little more about it, this would not work with dask='parallelized' which does not allow for chunking over core dimensions. You would have parallelize the function with dask yourself, e.g., with dask.array.map_blocks, but then you could use apply_ufunc with dask='allowed'.

One way to allow chunking across x would be to finish up dask.array.einsum: https://github.com/dask/dask/issues/732

I'm reluctant to add reduce_func to xarray because it isn't clear to me exactly what the underlying abstraction is. It's something like a gufunc, but does a little bit more. Also, ideally we'd like this to be in dask.array, maybe as part of dask.array.apply_gufunc (https://github.com/dask/dask/pull/3109).

For this specific problem, I think you could solve it with xarray.apply_ufunc by writing something like a gufunc that keeps the reduced axis as size 1 to apply to each chunk, and afterwards summing up along that dimension.

For two inputs, don't we use dask.array.tensordot?

Have you tried the new xarray.dot()? That might be even faster for this case.