diff --git a/src/fast_array_utils/_plugins/dask.py b/src/fast_array_utils/_plugins/dask.py index f9f93ca..535a29c 100644 --- a/src/fast_array_utils/_plugins/dask.py +++ b/src/fast_array_utils/_plugins/dask.py @@ -1,10 +1,7 @@ # SPDX-License-Identifier: MPL-2.0 from __future__ import annotations -import numpy as np - -# Other lookup candidates: tensordot_lookup and take_lookup -from dask.array.dispatch import concatenate_lookup +from dask.array.dispatch import concatenate_lookup, take_lookup, tensordot_lookup from scipy.sparse import sparray, spmatrix @@ -13,11 +10,12 @@ def patch() -> None: # pragma: no cover """Patch dask to support sparse arrays. - See + See """ # Avoid patch if already patched or upstream support has been added - if concatenate_lookup.dispatch(sparray) is not np.concatenate: + if concatenate_lookup.dispatch(sparray) is not concatenate_lookup.dispatch(spmatrix): return - concatenate = concatenate_lookup.dispatch(spmatrix) - concatenate_lookup.register(sparray, concatenate) + concatenate_lookup.register(sparray, concatenate_lookup.dispatch(spmatrix)) + tensordot_lookup.register(sparray, tensordot_lookup.dispatch(spmatrix)) + take_lookup.register(sparray, take_lookup.dispatch(spmatrix))