Improve Dask support

[krokosik]

The documentation clearly states that Dask is supported, yet all of the apply_ufunc do not set the dask prop, which defaults to "forbidden". I would propose a solution to have a global setting for this prop to not override default xarray behaviour. Some kind of singleton or even a global monkey patch, like with default dims.

[ahartikainen]

Are there some apply_ufunc calls where user can not give e.g. dask="allowed".

Working with dask is usually not too simple, so I'm not sure if enabling it by default makes sense.

[krokosik]

I'm pretty sure dask has almost none of the linalg methods natively supported, so parallelized would be preferable.

Passing this kwarg does not mean that everything gets converted to Dask, it just allows dask arrays to be passed into these calls. The default behaviour throws an error.

[OriolAbril]

Would it be ok if I convert this to a parent issue with multiple children? Or rename the issue and add related ones? I think there are several separated elements to this.

1. From my point of view, Dask is supported. The issue is it is only supported through function per function use of dask=xyz. If a function were not to pass the **kwargs to apply_ufunc to allow for this then we should also open dedicated issues to address the broken ones. e.g. if you do eig(dask_da, dims=..., dask="allowed/parallelized") it should already work
2. The dask support documentation page is a placeholder not an actual documentation page, it would be nice to add some actual examples in there.
3. I agree it would be interesting to have implicit/automatic support for Dask. I am not completely sure what is the best way forward though. It has been a few years already but I didn't find either allowed or parallelized to make much sense as global defaults. Even parallelized broke from time to time and when allowed can be used it is generally a much better alternative.

I think the discussion so far has been about point 3 only so we could focus this issue on this particular point and open an independent issue to discuss doc improvements. I am not sure what you have tried so far so I am not sure there are any functions that would require opening an issue for point 1 but I figured I'd mention it too anyway.

[krokosik]

Ohh right, sorry for the mistake, somehow I didn't connect that kwargs can be used for that. I will update the title.

Indeed, when I pass parallelized, eig just works, but like I said, not many linalg methods are actually native to Dask. It seems their docs do not render a list, but I managed to grab:
array scoped:

• einsum
• matmul
  linalg scoped:
• cholesky
• inv
• lstsq
• lu
• norm
• qr
• solve
• solve_triangular
• svd
• svd_compressed
• sfqr
• tsqr

A thorough search would reveal more probably.

With that being said, after passing applicable dask keywords, almost everything works out of the box. The issues I encountered are:

• inv throwing ValueError: Dimension must be 2 to perform lu decomposition with allowed
• inv being unable to infer output dtype with parallelized (works when I just set it to input dtype)
• same issues as above with solve
• trace throwing TypeError: trace() got an unexpected keyword argument 'out'. When I remove it or change to parallelized, it works
  Otherwise:
• eig and pinv work with parallelize
• matmul works with allowed

[krokosik]

@OriolAbril in my opinion, this package should leave the Dask support as is, just document it better. The reason for that is the spotty and not well documented coverage of Dask native linalg methods, meaning that supporting the allowed/parallelized keywords out of the box (withing xe package) will place a big maintenance burden on you, having to deal with different dask quirks across versions.

On the other hand, the end user is perfectly capable of picking correct keyword arguments themselves, despite what I originally wrote in the issue description ;). As such, an expanded documentation page about Dask support and limitations is most likely sufficient.

[OriolAbril]

Finally got to it, thanks for all the help @krokosik!

[krokosik]

Awesome, thank you!