issue_comments

Thanks to all reviewers! Great job!

@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.

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).

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.

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...

@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.

@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?

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.

@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.

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.

@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).

@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.

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 @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.

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.