don't import in the function call

dask-contrib · Dec 7, 2024 · 57a2472 · 57a2472
1 parent cee57c8
commit 57a2472
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/dask_awkward/lib/core.py b/src/dask_awkward/lib/core.py
@@ -61,6 +61,9 @@
     is_empty_slice,
 )
 
+if _dask_uses_tasks:
+    from dask._task_spec import TaskRef
+
 if TYPE_CHECKING:
     from awkward.contents.content import Content
     from awkward.forms.form import Form
@@ -1933,8 +1936,6 @@ def partitionwise_layer(
             numblocks[arg.name] = (1,)
         elif isinstance(arg, Delayed):
             if _dask_uses_tasks:
-                from dask._task_spec import TaskRef
-
                 pairs.extend([TaskRef(arg.key), None])
             else:
                 pairs.extend([arg.key, None])