typetracer fix

dask-contrib · Nov 15, 2023 · a776a4f · a776a4f
1 parent a1b55a2
commit a776a4f
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/dask_awkward/lib/structure.py b/src/dask_awkward/lib/structure.py
@@ -26,6 +26,7 @@
  DaskAwkwardNotImplemented,
  IncompatiblePartitions,
  borrow_docstring,
+ first,
 )
 
 if TYPE_CHECKING:
@@ -595,6 +596,9 @@ def nan_to_num(
 
 
 def _numaxis0(*integers):
+ f = first(integers)
+ if isinstance(f, TypeTracerArray):
+ return f
  return np.sum(np.array(integers))