fix: Finalize dask patch (#113)

Author: flying-sheep

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 <https://github.com/dask/dask/blob/4d71629d1f22ced0dd780919f22e70a642ec6753/dask/array/backends.py#L212-L232>
+    See <https://github.com/dask/dask/blob/d9b5c5b0256208f1befe94b26bfa8eaabcd0536d/dask/array/backends.py#L239-L241>
     """
     # 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))