Repository navigation
Improve Dask support #92
Description
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.
Activity
ahartikainen commentedon Apr 16, 2026
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 commentedon Apr 16, 2026
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 commentedon Apr 16, 2026
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.
dask=xyz. If a function were not to pass the**kwargstoapply_ufuncto allow for this then we should also open dedicated issues to address the broken ones. e.g. if you doeig(dask_da, dims=..., dask="allowed/parallelized")it should already workI 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 commentedon Apr 16, 2026
Ohh right, sorry for the mistake, somehow I didn't connect that
kwargscan be used for that. I will update the title.Indeed, when I pass
parallelized, eig just works, but like I said, not manylinalgmethods are actually native to Dask. It seems their docs do not render a list, but I managed to grab:array scoped:
linalg scoped:
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:
invthrowingValueError: Dimension must be 2 to perform lu decompositionwith allowedinvbeing unable to infer output dtype with parallelized (works when I just set it to input dtype)solvetracethrowingTypeError: trace() got an unexpected keyword argument 'out'. When I remove it or change to parallelized, it worksOtherwise:
eigandpinvwork with parallelizematmulworks with allowed[-]No Dask support[/-][+]Improve Dask support[/+]krokosik commentedon Jul 6, 2026
@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/parallelizedkeywords 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 commentedon Jul 20, 2026
Finally got to it, thanks for all the help @krokosik!
krokosik commentedon Jul 21, 2026
Awesome, thank you!