only check blockwise for Task; grab Task, TaskRef from blockwise

dask-contrib · Dec 17, 2024 · 64b2a66 · 64b2a66
1 parent 731ce1d
commit 64b2a66
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/dask_awkward/layers/layers.py b/src/dask_awkward/layers/layers.py
@@ -12,10 +12,10 @@
 
 from dask_awkward.utils import LazyInputsDict
 
-_dask_uses_tasks = hasattr(dask, "_task_spec") and hasattr(dask.blockwise, "Task")
+_dask_uses_tasks = hasattr(dask.blockwise, "Task")
 
 if _dask_uses_tasks:
-    from dask._task_spec import Task, TaskRef
+    from dask.blockwise import Task, TaskRef
 
 if TYPE_CHECKING:
     from awkward import Array as AwkwardArray