Thanks to all reviewers! Great job!

ok then - let's do this. Thanks a lot @kmuehlbauer

Hello @kmuehlbauer! Thanks for updating this PR. We checked the lines you've touched for PEP 8 issues, and found:

There are currently no PEP 8 issues detected in this Pull Request. Cheers! :beers:

Comment last updated at 2020-08-19 05:41:15 UTC

Nice! Unless @dcherian has any additional comments I'll merge in a few days

@mathause I've merged latest master into this PR to hopefully get all tests green. The former had some problems with a conda error in MinimumVersions job.

Please let me know, if there is anything for me to do, to get this merged.

While having a last review I've found another small glitch. I'll come back the next days to see, if anything needs to be done from reviewers side.

I think all of these can be done in a new PR, we just have to make sure to include them in the next release (which might need to be soon so we regain compatibility with the most recent pandas).

Great, than it looks like it's finally done. :smiley:

@keewis @mathause Thanks for the review. I've added a checklist in the first post with "open issues" with this PR, which might be solved in a follow up PR. Would be good to know which need to go in here, so I can add this.

@mathause I've merged with master ( + resolved conflicts) and addressed the latest review comments by @dcherian.

@dcherian Thanks for your latest review comments. I'll try to come up with a solution after merging with latest master (and resolving conflicts).

gentle ping @shoyer @dcherian I think this is not too far away

I had another look and seems fine to me. So let's ping @shoyer & @dcherian for a review.

How should we proceed with this PR?

@mathause This should now raise for both type of 'output_sizes' errors, wrong dim name in 'output_sizes' and missing dim name in 'output_sizes'. I'll catch up with any review tomorrow morning CET.

Could you use signature.all_output_core_dims ?

python if key not in signature.output_core_dims:

This does not work, because signature.output_core_dims is a list of tuples. So this would raise in any case. But I've no idea how to properly extract this into a single list/tuple...

You could go for something like this:

```python output_sizes_renamed = dict() for key, value in output_sizes.items():

if key not in signature.output_core_dims:
    raise ValueError("dimension name in 'output_sizes' must correspond to output_core_dims")
output_sizes_renamed[signature.dims_map[key]] = value

```

Instead of python output_sizes = { signature.dims_map[key]: value for key, value in output_sizes.items() }

@mathause It seems I can't find a nice way of checking/catching errors.

For now, this fails with KeyError inside xarray when output_sizes contains a wrong dimension name and this fails with KeyError inside dask, but with the not helpful dim0 message. Maybe this can be fixed in a follow up PR?

I would also leave combining of output_core_dims and output_sizes for another PR.

Ok, yes I am happy with erroring ourselves.

@mathause I've tested your suggestion but I do not know, what better error message you expect from the dask side.

• xr.apply_ufunc(..., output_core_dims=[["z"]], output_sizes={"z": 1}) signature.to_gufunc_string() will output ()->(dim0) (and output_sizes will be {'dim0': 1}
• xr.apply_ufunc(..., output_core_dims=[["z"]], output_sizes={"y": 1}) signature.to_gufunc_string() will output ()->(dim0) (and output_sizes will be {'y': 1}

So the second will fail with KeyError: 'dim0'.

I think we can't get proper error messages from dask side, because dask doesn't know anything about the new dimension name ( at least not from signature.to_gufunc_string()). Should we catch the error and add a meaningful output?

Yes, I think this got lost in a resolved comment.

output_core_dims=[{"abc": 2}] looks very elegant. But it is passed to _UFuncSignature(input_core_dims, output_core_dims) so you would probably need to take it apart again in the code (or adapt _UFuncSignature to handle this case). Not sure how involved this is. Could potentially be left for a future PR?

I made one more suggestion to catch the case

python xr.apply_ufunc(..., output_core_dims=[["abc"]], output_sizes={"xyz": 2})

Yes, I've got the email, but did not find it here. Thanks! Would it not be useful to combine output_core_dims and output_sizes, eg.

python xr.apply_ufunc(..., output_core_dims=[{"abc": 2]])

@mathause @shoyer @dcherian @max-sixty

This PR has gone a long way from first enabling multiple outputs to implementation of dask.apply_gufunc. Thanks for all your help, comments and suggestions. As xarray 0.16.0 is handled in the next days (or so I understood from reading the comments here #4031), I'm asking myself, if this can be finalized for 0.16.0?

I appreciate if we could work on any remaining problems to get this PR merged.

@dcherian @shoyer Only the upstream-dev is failing with the sparse-tests (#4146). This is ready for another review cycle. Please let me know, if there is more to work on this PR.

@mathause I reverted the mentioned error catching. I created a dedicated test for testing the dask vectorize without given dtype. Hope that fits your requirements.

@mathause Sorry for the delay. I have currently no other idea how to circumvent the use of dask_gufunc_kwargs also for the non-dask vectorize. I'm happy to adapt the code to a more convenient and readable solution.

This is ready for another review cycle. :smiley:

@mathause Thanks, I had just approached that particular line while debugging, when your message came in :grin:. I'll add an xfail for that dask version first.

Could be "Avoid adjusting gufunc’s meta dtype twice" in https://docs.dask.org/en/latest/changelog.html#id84

Not sure if it is better to xfail or bump the version of dask (which could be done ipython ci/min_deps_check.py ci/requirements/py36-min-all-deps.yml)

@mathause Any idea why the meta test fails in Linux py36-min-all-deps?

@mathause I've added your tests from #4022 to this PR (dtype/meta).

OK, vectorization has to be inside apply_variable_ufunc to be applied properly on dask and non-dask configurations.

For the time being I use dask_gufunc_kwargs to distribute vectorize and output_dtypes into apply_variable_ufunc. Maybe we have to decide on more sensible name for that dict.

Currently tests are failing for some test_missing. Any hints appreciated.

Thanks @shoyer and @mathause. I've tried to adapt the PR to your comments.

• The meta and output_sizes parameters should be and the vectorize and output_dtypes parameters can be given within the dask_gufunc_kwargs (which have priority). This would also be the place for allow_rechunk. Everything is mentioned in the docstring and warnings are issued to use the new dask_gufunc_kwargs. The dask_gufunc_kwargs itself is used to hold all keyword arguments for dask.array.apply_gufunc.

• vectorization is done inside apply_ufunc only for dask not parallelized. In that case it is deferred to dask.array.apply_gufunc.

• for chunking problems (core dimension with multiple chunks and non-core dimensions with different chunksizes) we let dask raise the errors.

• un-equal chunking along non-core dimensions

The second is not so trivial. I see three possibilities (1) just error, (2) try dask.array.apply_gufunc if that fails issue a warning and use the old apply_blockwise (3) figure out ourselves if non-core dimensions (called loop dimensions in dask) are not-equally chunked, issue a warning and re-chunk them ourselves. Maybe @shoyer and @dcherian can weight in here.

I am pretty confident that existing behavior of xarray.apply_ufunc with un-equal chunks along non-core dimensions is entirely broken. I am OK with just erroring for now.

Looking at the internals of dask.array.apply_gufunc, it looks like it passes on **kwargs to blockwise. blockwise gets called twice, first with only **kwargs and then with both meta and **kwargs: https://github.com/dask/dask/blob/77628f2d5248cc61cc366cac3e400b6df5c654c1/dask/array/gufunc.py#L422 https://github.com/dask/dask/blob/77628f2d5248cc61cc366cac3e400b6df5c654c1/dask/array/gufunc.py#L439

So it seems we can actually still pass on a meta argument explicitly, at least in the first case.

Here's my suggestion for moving forwarding:

• If meta was explicitly set (meta is not None), pass it into apply_gufunc in **kwargs
• Otherwise, omit meta.

No need to worry about the second one:

https://github.com/pydata/xarray/blob/e5cc19cd8f8a69e0743f230f5bf51b7778a0ff96/xarray/core/computation.py#L704

What I was trying to say - this is a change that can break existing code and it is not easy to do a proper deprecation cycle (i.e. only throw a warning). (I think allow_rechunk=True is not a good default.)

Yes, sure, I didn't think about the effects on existing code. We might need to capture both ValueErrors and handle differently. I'll try to come up with a proposal for this.

output_dtypes: yes I think the docs should be updated

What I was trying to say - this is a change that can break existing code and it is not easy to do a proper deprecation cycle (i.e. only throw a warning). (I think allow_rechunk=True is not a good default.)

Ok an easier idea (could be brittle, though):

```python try: res = da.apply_gufunc(...)

catch error that was not raised in earlier versions of xr.apply_ufunc

except ValueError as e: if "with different chunksize present" in str(e): warnings.warn(str(e) + " This will throw an error in the future.") res = da.apply_gufunc(..., allow_rechunk=True) else: raise

``` However, I am only throwing around ideas - this is certainly not up to me. Let's see what the others think. Maybe they have an opinion or a better idea.

@mathause Thanks for the elaborate review. But, I'm unsure how to move on with this.

I'll reply to some of your comments below.

Note that output_dtypes is used in vectorize so this may not be a candidate for dask_gufunc_kwargs.

Hmm, the docstring indicates that it is only used with dask='parallelized'. But in the code it can be used in any case. This should be clarified.

The second is not so trivial. I see three possibilities (1) just error, (2) try dask.array.apply_gufunc if that fails issue a warning and use the old apply_blockwise (3) figure out ourselves if non-core dimensions (called loop dimensions in dask) are not-equally chunked, issue a warning and re-chunk them ourselves. Maybe @shoyer and @dcherian can weight in here.

When looking at the code of dask.array.apply_gufunc it seems, that any possible re-chunking is done in dask.array.blockwise. I also do not see any special treatment of warning or re-chunking of non-core dimensions in the current _apply_blockwise. So from my reasoning we should be safe to use (4) try dask.array.apply_gufunc since dask.array.apply_gufunc will issue the appropriate ValueError if something is wrong with the chunking. (Only thing is that the ValueError concerning the non-core dimension doesn't mention allow_rechunk). If the ValueError is issued, the user can the just set allow_rechunk=True keyword.

although only two keywords seem to be worth feeding through (keepdims, allow_rechunk).

Just noticed, that output_dtypes and output_sizes could also be fed as gufunc_kwargs. So this would really simplify the calling signature of apply_ufunc if those kwargs could be handled that way.

I'll only be able to look at it properly next week.

@mathause I'll leave the PR unchanged and catch up with you next week.

@shoyer @dcherian Thanks for your comments. Please let me know, which tests should be added to check for any possible surprises with this change to apply_gufunc.

The problem is that we don't have any way to detect ahead of time whether the applied function already supports dask arrays (e.g., if it is built-up out of functions from dask.array). If it does, we don't want to set dask='parallelized' but rather let the function handle dask arrays itself.

(Att: no native english speaker here, so bear with me, if something sounds clunky or not exactly matching) Then we would have to keep the dask='forbidden' as default, as well as parallelized and allowed to force the decision to the user. Maybe the keyword settings itself could be a bit more clear. In the allowed-case the function in question has to natively support dask-arrays. So I would use dask='native' in that case. For the parallelized-case this PR proposes to use dask.array.apply_gufunc (generalized ufunc). So either we stick to parallelized or we try to find a better fitting name (eg. dask='gufunc').

For the keywords I think @dcherian 's proposal of something like dask_gufunc_kwargs (or gufunc_kwargs) is useful (would match with dask='gufunc'), although only two keywords seem to be worth feeding through (keepdims, allow_rechunk).

good point @mathause. Looks like apply_gufunc tries to make blockwise infer meta and then does its own thing when that fails: https://github.com/dask/dask/blob/3b92efff2e779f59e95e05af9b8f371d56227d02/dask/array/gufunc.py#L417-L440 . I don't understand how this can work in all possible cases.

I don't see meta listed in the docs which is why it thought it's not needed. But if it is handled in dask.array.apply_gufunc it can of course stay.

I only realised the exact distinction between "allowed" and "parallelized" today - i.e. that "parallelized" is kind of the dask equivalent of np.vectorize. I can suggest something for the docstring (e.g. prefer "allowed" if ``func`` natively handles dask arrays or so)

If it does, we don't want to set dask='parallelized' but rather let the function handle dask arrays itself.

I think we still need all the current options for the dask kwarg. There can be disastrous consequences so it's good to make users explicitly choose the behaviour they want.

howto properly handle deprecation (eg. meta)

I don't think we should deprecate meta. Not all user functions can deal with zero shaped inputs, so automatically inferring meta need not always work. We've had to add a similar feature for map_blocks (#3575) so I think meta should stay.

keywords of dask.array.apply_gufunc

Shall we add a new dask_gufunc_kwargs and pass that down to appy_gufunc?

Maybe this is too defensive/surprising, and could be relaxed.

You would remove the daks="forbidden" branch and not the dask="parallelized"?

For the functions that don't handle dask arrays gracefully, dask="parallelized" would be the better option?

This is probably another good motivation: defaulting to dask='forbidden' forces users to make an explicit choice about whether or not use dask='parallelized'.

The problem is that we don't have any way to detect ahead of time whether the applied function already supports dask arrays (e.g., if it is built-up out of functions from dask.array). If it does, we don't want to set dask='parallelized' but rather let the function handle dask arrays itself.

Maybe this is too defensive/surprising, and could be relaxed.

You would remove the daks="forbidden" branch and not the dask="parallelized"?

For the functions that don't handle dask arrays gracefully, dask="parallelized" would be the better option?

Very cool - good progress.

• I guess you'll have to properly deprecate meta, something along the lines of: meta is no longer necessary and has no effect. it will be removed in a future version
• I think it would be good to pass allow_rechunk.

I'll only be able to look at it properly next week.

The original motivation for requiring dask='allowed' is that I was concerned that users would put a function that coerces its arguments into NumPy arrays into apply_ufunc (e.g., like many functions from SciPy), which could have surprisingly bad performance when called on dask arrays due to automatic coercion.

Maybe this is too defensive/surprising, and could be relaxed. We don't really have any guard-rails like this elsewhere in xarray.

@mathause All tests green, good starting point for review. Please notify other people who should have a look at this.

There are still things to discuss: - keywords of dask.array.apply_gufunc - howto properly handle deprecation (eg. meta) - doscstring - adding/revision of tests

I've given this a try, but this will need some design decisions.

• currently vectorize is handled in any case if requested, before falling through the if/else. dask.array.apply_gufunc takes vectorize as parameter, so for dask we do not need to apply vectorization. We would need to apply vectorize only for non-dask cases (maybe just before calling the final function).
• currently dask is only handled, if dask keyword is issued ('allowed' and parallelized). From my perspective the dask keyword is not needed any more. We could just divert to the apply_gufunc when dask backed arrays are detected.

This will really have much impact on the code/tests. I'll come up with a updated PullRequest in short time, but any thoughts /remarks whatsoever are very much appreciated.

@mathause Only one breaking test which is connected to the meta stuff you handle in #4022. Any suggestions on that topic? I've removed meta completely since it was only needed in dask.array.blockwise but not in dask.array.apply_gufunc.

@mathause Great! Seems like this works better, thanks! Will update the PR after some more tests etc.

Should you keep the dask="allowed" branch? That might solve some of the issues (instead of using allow_rechunk=True) . For example dask has it's own einsum implementation (used in xr.dot) so it may not be necessary to pipe that through dask.array.apply_gufunc (from the docstring: this function is like np.vectorize, but for the blocks of dask arrays). However, I am only speculating.

From looking at the tests, it seems that setting allow_rechunk=True would solve many issues. I've no idea about the implications on memory usage (docstring: "Warning: enabling this can increase memory usage significantly"). Might apply_gufunc not be suitable for processing of eg. dot?

@mathause @shoyer

First serve of trying to use dask.array.apply_gufunc in xr.apply_ufunc. I've added a list with problems in the topmost comment of this PR, to not loose track of this. Please enhance that list, if needed.

Most problematic issue now: xr.dot doesn't work well with apply_gufunc with regard to core dimensions and chunking.

Thanks! That would be quite cool!

dask.array.apply_gufunc does some fancy stuff with np.vectorize (to determine the output_dtypes) so I would not vectorize the function ourselves. For the other I don't have a qualified opinion.

• Would it be possible to replace the call to dask.array.blockwise (for one output variable) with dask.array.apply_gufunc? Do you know why blockwise is used further below and not dask.array.apply_gufunc? I assume it's due to historical reasons but I am not sure.

AFAIK, apply_gufunc wasn't available at the time these functions were introduced. Good chance, that apply_gufunc can be used for handling single output dask too.

Exactly. It would be nice remove the use of blockwise entirely in favor of apply_gufunc.

Ah yes I see (https://github.com/pydata/xarray/issues/1815#issuecomment-440089606). dask.array.apply_gufunc should also be able to handle one output only.

A complete refactor of apply_ufunc would be quite some challenge.

Indeed - I think it could simplify _apply_blockwise (and might make the meta keword obsolete) but it would be good if someone with more experience of dask could weigh in.

@dcherian @shoyer

Thanks @mathause for your comments and raising those questions. JFTR, I was taking the road from #1815, so my explicit use-case was the multiple (dask) outputs.

• It might be good to add a test with a reduction and one with vectorize=True.

I'll try to add some tests for the multiple output using dask.

• Would it be possible to replace the call to dask.array.blockwise (for one output variable) with dask.array.apply_gufunc? Do you know why blockwise is used further below and not dask.array.apply_gufunc? I assume it's due to historical reasons but I am not sure.

AFAIK, apply_gufunc wasn't available at the time these functions were introduced. Good chance, that apply_gufunc can be used for handling single output dask too.

• dask.array.apply_gufunc does all sorts of stuff - e.g. infer meta. This could potentially solve #4015 (pull #4022) and simplify the call signature of apply_ufunc?

That's a good question. If you want me to go the long way, please be aware, that I'm a novice in xarray as well as in dask. A complete refactor of apply_ufunc would be quite some challenge.

• It might be good to add a test with a reduction and one with vectorize=True.
• Would it be possible to replace the call to dask.array.blockwise (for one output variable) with dask.array.apply_gufunc? Do you know why blockwise is used further below and not dask.array.apply_gufunc? I assume it's due to historical reasons but I am not sure.
• dask.array.apply_gufunc does all sorts of stuff - e.g. infer meta. This could potentially solve #4015 (pull #4022) and simplify the call signature of apply_ufunc?

https://github.com/dask/dask/blob/3573b2ddca81aeb41a7def6dd4194020f853ab18/dask/array/gufunc.py#L175

This is ready for review from my side.

This would need some docstring changing too. But I first want to check, if I've missed anything vital in the implementation.