Matrix multiplication inconsistent for dask and sparse: dask @ sparse works, sparse @ dask fails

[brendan-m-murphy]

What is your issue?

The order of matrix multiplication matters for dask and sparse arrays, but probably shouldn't.

Here is an example:

import numpy as np
from sparse import COO
import xarray as xr

sparse_da = xr.DataArray(COO.from_numpy(np.arange(10)))
dask_da = xr.DataArray(np.arange(100).reshape(10, 10).chunk({"dim_1": 5})

# this works:
dask_da @ sparse_da

# this raises a TypeError for "unsupported types"
sparse_da @ dask_da

# this works as expected
sparse_da @ dask_da.as_numpy()

In the workflow where this is used, the dask array has no chunks along its common dimensions with the sparse array, so it seems like sparse @ chunk should be fine. Also, in this workflow, loading the dask array into memory or making the sparse array dense would use a very large amount of memory.

[keewis]

this is a side-effect of using opt_einsum: while

sparse_da @ dask_da

fails, this:

with xr.set_options(use_opt_einsum=False):
    sparse_da @ dask_da

works (but is slower). You could also avoid this issue by converting sparse_da to a single-chunk dask array using .chunk().

(not sure what to do to fix this in xarray, but I agree that this is annoying)

[brendan-m-murphy]

Thanks for the suggestions!

This is the ufunc we've been using instead: https://github.com/openghg/openghg_inversions/blob/sparse-xarray-fix/openghg_inversions/array_ops.py#L75. opt_einsum ends up calling sparse.tensordot, so I made this ufunc just call sparse.tensordot immediately. There's a bit of a "gotcha" where sparse.tensordot leaves extra broadcast dimensions around, so they need to be removed.

I might just rewrite it to apply @ in the opposite order, or use your chunking suggestion. (Or remove the function and just inline this code.)

Maybe xarray could check for the case of sparse and dask and either swap the order or convert the sparse array into a single chunk dask array? I guess opt-einsum could probably fix this too, but it seems like they just check if sparse has a tensordot attribute and apply it, so they would also need to insert some ugly logic to check both operands and decide what tensordot to use.

[dcherian]

Can you open an issue at pydata/sparse instead?

[brendan-m-murphy]

Can you open an issue at pydata/sparse instead?

Sure, I can see what they say. I think they would need to change tensordot to accept dask arrays (or any duck array?). At least in this case, the non-sparse array just needs to support indexing (I think).