Repository: cubed-dev/cubed
Pull request #686 (2025-01-30 12:47:06)

Add scalar support to binary element-wise ops
Part of #678 


Patch:
diff --git a/cubed/array_api/dtypes.py b/cubed/array_api/dtypes.py
index 19af81ad..937172c4 100644
--- a/cubed/array_api/dtypes.py
+++ b/cubed/array_api/dtypes.py
@@ -90,7 +90,9 @@ _dtype_categories = {
 
 
 # A Cubed-specific utility.
-def _upcast_integral_dtypes(x, dtype=None, *, allowed_dtypes=("numeric",), fname=None, device=None):
+def _upcast_integral_dtypes(
+    x, dtype=None, *, allowed_dtypes=("numeric",), fname=None, device=None
+):
     """Ensure the input dtype is allowed. If it's None, provide a good default dtype."""
     dtypes = __array_namespace_info__().default_dtypes(device=device)
 
@@ -116,3 +118,16 @@ def _upcast_integral_dtypes(x, dtype=None, *, allowed_dtypes=("numeric",), fname
             dtype = x.dtype
 
     return dtype
+
+
+def _promote_scalars(x1, x2, op):
+    """Promote at most one of x1 or x2 to an array from a Python scalar"""
+    x1_is_scalar = isinstance(x1, (int, float, complex, bool))
+    x2_is_scalar = isinstance(x2, (int, float, complex, bool))
+    if x1_is_scalar and x2_is_scalar:
+        raise TypeError(f"At least one of x1 and x2 must be an array in {op}")
+    elif x1_is_scalar:
+        x1 = x2._promote_scalar(x1)
+    elif x2_is_scalar:
+        x2 = x1._promote_scalar(x2)
+    return x1, x2
diff --git a/cubed/array_api/elementwise_functions.py b/cubed/array_api/elementwise_functions.py
index 217a6788..11572f49 100644
--- a/cubed/array_api/elementwise_functions.py
+++ b/cubed/array_api/elementwise_functions.py
@@ -8,6 +8,7 @@ from cubed.array_api.dtypes import (
     _integer_dtypes,
     _integer_or_boolean_dtypes,
     _numeric_dtypes,
+    _promote_scalars,
     _real_floating_dtypes,
     _real_numeric_dtypes,
     complex64,
@@ -44,6 +45,7 @@ def acosh(x, /):
 
 
 def add(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "add")
     if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in add")
     return elemwise(nxp.add, x1, x2, dtype=result_type(x1, x2))
@@ -68,6 +70,7 @@ def atan(x, /):
 
 
 def atan2(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "atan2")
     if x1.dtype not in _real_floating_dtypes or x2.dtype not in _real_floating_dtypes:
         raise TypeError("Only real floating-point dtypes are allowed in atan2")
     return elemwise(nxp.atan2, x1, x2, dtype=result_type(x1, x2))
@@ -80,6 +83,7 @@ def atanh(x, /):
 
 
 def bitwise_and(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "bitwise_and")
     if (
         x1.dtype not in _integer_or_boolean_dtypes
         or x2.dtype not in _integer_or_boolean_dtypes
@@ -95,12 +99,14 @@ def bitwise_invert(x, /):
 
 
 def bitwise_left_shift(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "bitwise_left_shift")
     if x1.dtype not in _integer_dtypes or x2.dtype not in _integer_dtypes:
         raise TypeError("Only integer dtypes are allowed in bitwise_left_shift")
     return elemwise(nxp.bitwise_left_shift, x1, x2, dtype=result_type(x1, x2))
 
 
 def bitwise_or(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "bitwise_or")
     if (
         x1.dtype not in _integer_or_boolean_dtypes
         or x2.dtype not in _integer_or_boolean_dtypes
@@ -110,12 +116,14 @@ def bitwise_or(x1, x2, /):
 
 
 def bitwise_right_shift(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "bitwise_right_shift")
     if x1.dtype not in _integer_dtypes or x2.dtype not in _integer_dtypes:
         raise TypeError("Only integer dtypes are allowed in bitwise_right_shift")
     return elemwise(nxp.bitwise_right_shift, x1, x2, dtype=result_type(x1, x2))
 
 
 def bitwise_xor(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "bitwise_xor")
     if (
         x1.dtype not in _integer_or_boolean_dtypes
         or x2.dtype not in _integer_or_boolean_dtypes
@@ -172,6 +180,7 @@ def conj(x, /):
 
 
 def copysign(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "copysign")
     if x1.dtype not in _real_numeric_dtypes or x2.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in copysign")
     return elemwise(nxp.copysign, x1, x2, dtype=result_type(x1, x2))
@@ -190,6 +199,7 @@ def cosh(x, /):
 
 
 def divide(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "divide")
     if x1.dtype not in _floating_dtypes or x2.dtype not in _floating_dtypes:
         raise TypeError("Only floating-point dtypes are allowed in divide")
     return elemwise(nxp.divide, x1, x2, dtype=result_type(x1, x2))
@@ -208,6 +218,7 @@ def expm1(x, /):
 
 
 def equal(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "equal")
     return elemwise(nxp.equal, x1, x2, dtype=nxp.bool)
 
 
@@ -221,20 +232,24 @@ def floor(x, /):
 
 
 def floor_divide(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "floor_divide")
     if x1.dtype not in _real_numeric_dtypes or x2.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in floor_divide")
     return elemwise(nxp.floor_divide, x1, x2, dtype=result_type(x1, x2))
 
 
 def greater(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "greater")
     return elemwise(nxp.greater, x1, x2, dtype=nxp.bool)
 
 
 def greater_equal(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "greater_equal")
     return elemwise(nxp.greater_equal, x1, x2, dtype=nxp.bool)
 
 
 def hypot(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "hypot")
     if x1.dtype not in _real_numeric_dtypes or x2.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in hypot")
     return elemwise(nxp.hypot, x1, x2, dtype=result_type(x1, x2))
@@ -269,10 +284,12 @@ def isnan(x, /):
 
 
 def less(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "less")
     return elemwise(nxp.less, x1, x2, dtype=nxp.bool)
 
 
 def less_equal(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "less_equal")
     return elemwise(nxp.less_equal, x1, x2, dtype=nxp.bool)
 
 
@@ -301,12 +318,14 @@ def log10(x, /):
 
 
 def logaddexp(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "logaddexp")
     if x1.dtype not in _real_floating_dtypes or x2.dtype not in _real_floating_dtypes:
         raise TypeError("Only real floating-point dtypes are allowed in logaddexp")
     return elemwise(nxp.logaddexp, x1, x2, dtype=result_type(x1, x2))
 
 
 def logical_and(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "logical_and")
     if x1.dtype not in _boolean_dtypes or x2.dtype not in _boolean_dtypes:
         raise TypeError("Only boolean dtypes are allowed in logical_and")
     return elemwise(nxp.logical_and, x1, x2, dtype=nxp.bool)
@@ -319,30 +338,35 @@ def logical_not(x, /):
 
 
 def logical_or(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "logical_or")
     if x1.dtype not in _boolean_dtypes or x2.dtype not in _boolean_dtypes:
         raise TypeError("Only boolean dtypes are allowed in logical_or")
     return elemwise(nxp.logical_or, x1, x2, dtype=nxp.bool)
 
 
 def logical_xor(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "logical_xor")
     if x1.dtype not in _boolean_dtypes or x2.dtype not in _boolean_dtypes:
         raise TypeError("Only boolean dtypes are allowed in logical_xor")
     return elemwise(nxp.logical_xor, x1, x2, dtype=nxp.bool)
 
 
 def maximum(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "maximum")
     if x1.dtype not in _real_numeric_dtypes or x2.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in maximum")
     return elemwise(nxp.maximum, x1, x2, dtype=result_type(x1, x2))
 
 
 def minimum(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "minimum")
     if x1.dtype not in _real_numeric_dtypes or x2.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in minimum")
     return elemwise(nxp.minimum, x1, x2, dtype=result_type(x1, x2))
 
 
 def multiply(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "multiply")
     if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in multiply")
     return elemwise(nxp.multiply, x1, x2, dtype=result_type(x1, x2))
@@ -355,6 +379,7 @@ def negative(x, /):
 
 
 def not_equal(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "not_equal")
     return elemwise(nxp.not_equal, x1, x2, dtype=nxp.bool)
 
 
@@ -365,6 +390,7 @@ def positive(x, /):
 
 
 def pow(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "pow")
     if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in pow")
     return elemwise(nxp.pow, x1, x2, dtype=result_type(x1, x2))
@@ -381,6 +407,7 @@ def real(x, /):
 
 
 def remainder(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "remainder")
     if x1.dtype not in _real_numeric_dtypes or x2.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in remainder")
     return elemwise(nxp.remainder, x1, x2, dtype=result_type(x1, x2))
@@ -429,6 +456,7 @@ def square(x, /):
 
 
 def subtract(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "subtract")
     if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in subtract")
     return elemwise(nxp.subtract, x1, x2, dtype=result_type(x1, x2))
diff --git a/cubed/array_api/searching_functions.py b/cubed/array_api/searching_functions.py
index eeec8411..06a250ac 100644
--- a/cubed/array_api/searching_functions.py
+++ b/cubed/array_api/searching_functions.py
@@ -1,6 +1,6 @@
 from cubed.array_api.creation_functions import asarray, zeros_like
 from cubed.array_api.data_type_functions import result_type
-from cubed.array_api.dtypes import _real_numeric_dtypes
+from cubed.array_api.dtypes import _promote_scalars, _real_numeric_dtypes
 from cubed.array_api.manipulation_functions import reshape
 from cubed.array_api.statistical_functions import max
 from cubed.backend_array_api import namespace as nxp
@@ -88,13 +88,6 @@ def _searchsorted(x, y, side):
 
 
 def where(condition, x1, x2, /):
-    x1_is_scalar = isinstance(x1, (int, float, complex, bool))
-    x2_is_scalar = isinstance(x2, (int, float, complex, bool))
-    if x1_is_scalar and x2_is_scalar:
-        raise TypeError("At least one of x1 and x2 must be an array in where")
-    elif x1_is_scalar:
-        x1 = x2._promote_scalar(x1)
-    elif x2_is_scalar:
-        x2 = x1._promote_scalar(x2)
+    x1, x2 = _promote_scalars(x1, x2, "where")
     dtype = result_type(x1, x2)
     return elemwise(nxp.where, condition, x1, x2, dtype=dtype)
diff --git a/cubed/runtime/executors/asyncio.py b/cubed/runtime/asyncio.py
similarity index 60%
rename from cubed/runtime/executors/asyncio.py
rename to cubed/runtime/asyncio.py
index 6e4f7a2c..e0dfd266 100644
--- a/cubed/runtime/executors/asyncio.py
+++ b/cubed/runtime/asyncio.py
@@ -2,10 +2,30 @@ import asyncio
 import copy
 import time
 from asyncio import Future
-from typing import Any, AsyncIterator, Callable, Dict, Iterable, List, Optional, Tuple
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+)
+
+from aiostream import stream
+from aiostream.core import Stream
+from networkx import MultiDiGraph
 
 from cubed.runtime.backup import should_launch_backup
-from cubed.runtime.utils import batched
+from cubed.runtime.pipeline import visit_node_generations, visit_nodes
+from cubed.runtime.types import Callback, CubedPipeline
+from cubed.runtime.utils import (
+    batched,
+    handle_callbacks,
+    handle_operation_start_callbacks,
+)
 
 
 async def async_map_unordered(
@@ -100,3 +120,61 @@ async def async_map_unordered(
                 pending.update(new_tasks.keys())
                 t = time.monotonic()
                 start_times = {f: t for f in new_tasks.keys()}
+
+
+async def async_map_dag(
+    create_futures_func: Callable,
+    dag: MultiDiGraph,
+    callbacks: Optional[Sequence[Callback]] = None,
+    resume: Optional[bool] = None,
+    compute_arrays_in_parallel: Optional[bool] = None,
+    **kwargs,
+) -> None:
+    """
+    Asynchronous parallel map over multiple pipelines from a DAG, with support for backups and batching.
+    """
+    if not compute_arrays_in_parallel:
+        # run one pipeline at a time
+        for name, node in visit_nodes(dag, resume=resume):
+            handle_operation_start_callbacks(callbacks, name)
+            st = pipeline_to_stream(
+                create_futures_func, name, node["pipeline"], **kwargs
+            )
+            async with st.stream() as streamer:
+                async for result, stats in streamer:
+                    handle_callbacks(callbacks, result, stats)
+    else:
+        for gen in visit_node_generations(dag, resume=resume):
+            # run pipelines in the same topological generation in parallel by merging their streams
+            streams = [
+                pipeline_to_stream(
+                    create_futures_func, name, node["pipeline"], **kwargs
+                )
+                for name, node in gen
+            ]
+            merged_stream = stream.merge(*streams)
+            async with merged_stream.stream() as streamer:
+                async for result, stats in streamer:
+                    handle_callbacks(callbacks, result, stats)
+
+
+def pipeline_to_stream(
+    create_futures_func: Callable,
+    name: str,
+    pipeline: CubedPipeline,
+    **kwargs,
+) -> Stream:
+    """
+    Turn a pipeline into an asynchronous stream of results.
+    """
+    return stream.iterate(
+        async_map_unordered(
+            create_futures_func,
+            pipeline.mappable,
+            return_stats=True,
+            name=name,
+            func=pipeline.function,
+            config=pipeline.config,
+            **kwargs,
+        )
+    )
diff --git a/cubed/runtime/executors/dask.py b/cubed/runtime/executors/dask.py
index 8728384f..041fc099 100644
--- a/cubed/runtime/executors/dask.py
+++ b/cubed/runtime/executors/dask.py
@@ -1,33 +1,13 @@
 import asyncio
-from typing import (
-    Any,
-    AsyncIterator,
-    Callable,
-    Dict,
-    Iterable,
-    List,
-    Optional,
-    Sequence,
-    Tuple,
-    Union,
-)
-
-from aiostream import stream
-from aiostream.core import Stream
+from typing import Any, Callable, Dict, Optional, Sequence
+
 from dask.distributed import Client
 from networkx import MultiDiGraph
 
+from cubed.runtime.asyncio import async_map_dag
 from cubed.runtime.backup import use_backups_default
-from cubed.runtime.executors.asyncio import async_map_unordered
-from cubed.runtime.pipeline import visit_node_generations, visit_nodes
-from cubed.runtime.types import Callback, CubedPipeline, DagExecutor
-from cubed.runtime.utils import (
-    asyncio_run,
-    execution_stats,
-    gensym,
-    handle_callbacks,
-    handle_operation_start_callbacks,
-)
+from cubed.runtime.types import Callback, DagExecutor
+from cubed.runtime.utils import asyncio_run, execution_stats, gensym
 from cubed.spec import Spec
 
 
@@ -38,68 +18,6 @@ def run_func(input, pipeline_func=None, config=None, name=None, compute_id=None)
     return result
 
 
-async def map_unordered(
-    client: Client,
-    map_function: Callable[..., Any],
-    map_iterdata: Iterable[Union[List[Any], Tuple[Any, ...], Dict[str, Any]]],
-    retries: int = 2,
-    use_backups: bool = False,
-    batch_size: Optional[int] = None,
-    return_stats: bool = False,
-    name: Optional[str] = None,
-    **kwargs,
-) -> AsyncIterator[Any]:
-    def create_futures_func(input, **kwargs):
-        input = list(input)  # dask expects a sequence (it calls `len` on it)
-        key = name or gensym("map")
-        key = key.replace("-", "_")  # otherwise array number is not shown on dashboard
-        return [
-            (i, asyncio.ensure_future(f))
-            for i, f in zip(
-                input,
-                client.map(map_function, input, key=key, retries=retries, **kwargs),
-            )
-        ]
-
-    def create_backup_futures_func(input, **kwargs):
-        input = list(input)  # dask expects a sequence (it calls `len` on it)
-        key = name or gensym("backup")
-        key = key.replace("-", "_")  # otherwise array number is not shown on dashboard
-        return [
-            (i, asyncio.ensure_future(f))
-            for i, f in zip(input, client.map(map_function, input, key=key, **kwargs))
-        ]
-
-    async for result in async_map_unordered(
-        create_futures_func,
-        map_iterdata,
-        use_backups=use_backups,
-        create_backup_futures_func=create_backup_futures_func,
-        batch_size=batch_size,
-        return_stats=return_stats,
-        name=name,
-        **kwargs,
-    ):
-        yield result
-
-
-def pipeline_to_stream(
-    client: Client, name: str, pipeline: CubedPipeline, **kwargs
-) -> Stream:
-    return stream.iterate(
-        map_unordered(
-            client,
-            run_func,
-            pipeline.mappable,
-            return_stats=True,
-            name=name,
-            pipeline_func=pipeline.function,
-            config=pipeline.config,
-            **kwargs,
-        )
-    )
-
-
 def check_runtime_memory(spec, client):
     allowed_mem = spec.allowed_mem if spec is not None else None
     scheduler_info = client.scheduler_info()
@@ -114,40 +32,27 @@ def check_runtime_memory(spec, client):
             )
 
 
-async def async_execute_dag(
-    dag: MultiDiGraph,
-    callbacks: Optional[Sequence[Callback]] = None,
-    resume: Optional[bool] = None,
-    spec: Optional[Spec] = None,
-    compute_arrays_in_parallel: Optional[bool] = None,
-    compute_kwargs: Optional[Dict[str, Any]] = None,
-    **kwargs,
-) -> None:
-    compute_kwargs = compute_kwargs or {}
-    async with Client(asynchronous=True, **compute_kwargs) as client:
-        if spec is not None:
-            check_runtime_memory(spec, client)
-            if "use_backups" not in kwargs and use_backups_default(spec):
-                kwargs["use_backups"] = True
-        if not compute_arrays_in_parallel:
-            # run one pipeline at a time
-            for name, node in visit_nodes(dag, resume=resume):
-                handle_operation_start_callbacks(callbacks, name)
-                st = pipeline_to_stream(client, name, node["pipeline"], **kwargs)
-                async with st.stream() as streamer:
-                    async for result, stats in streamer:
-                        handle_callbacks(callbacks, result, stats)
-        else:
-            for gen in visit_node_generations(dag, resume=resume):
-                # run pipelines in the same topological generation in parallel by merging their streams
-                streams = [
-                    pipeline_to_stream(client, name, node["pipeline"], **kwargs)
-                    for name, node in gen
-                ]
-                merged_stream = stream.merge(*streams)
-                async with merged_stream.stream() as streamer:
-                    async for result, stats in streamer:
-                        handle_callbacks(callbacks, result, stats)
+def dask_create_futures_func(
+    client,
+    function: Callable[..., Any],
+    name: Optional[str] = None,
+    retries: Optional[str] = None,
+):
+    def create_futures_func(input, **kwargs):
+        input = list(input)  # dask expects a sequence (it calls `len` on it)
+        key = name or gensym("map")
+        key = key.replace("-", "_")  # otherwise array number is not shown on dashboard
+        if "func" in kwargs:
+            kwargs["pipeline_func"] = kwargs.pop("func")  # rename to avoid clash
+        return [
+            (i, asyncio.ensure_future(f))
+            for i, f in zip(
+                input,
+                client.map(function, input, key=key, retries=retries, **kwargs),
+            )
+        ]
+
+    return create_futures_func
 
 
 class DaskExecutor(DagExecutor):
@@ -172,7 +77,7 @@ class DaskExecutor(DagExecutor):
     ) -> None:
         merged_kwargs = {**self.kwargs, **kwargs}
         asyncio_run(
-            async_execute_dag(
+            self._async_execute_dag(
                 dag,
                 callbacks=callbacks,
                 resume=resume,
@@ -182,3 +87,39 @@ class DaskExecutor(DagExecutor):
                 **merged_kwargs,
             )
         )
+
+    async def _async_execute_dag(
+        self,
+        dag: MultiDiGraph,
+        callbacks: Optional[Sequence[Callback]] = None,
+        resume: Optional[bool] = None,
+        spec: Optional[Spec] = None,
+        compute_kwargs: Optional[Dict[str, Any]] = None,
+        compute_arrays_in_parallel: Optional[bool] = None,
+        **kwargs,
+    ) -> None:
+        compute_kwargs = compute_kwargs or {}
+        retries = kwargs.pop("retries", 2)
+        name = kwargs.get("name", None)
+        async with Client(asynchronous=True, **compute_kwargs) as client:
+            if spec is not None:
+                check_runtime_memory(spec, client)
+                if "use_backups" not in kwargs and use_backups_default(spec):
+                    kwargs["use_backups"] = True
+
+            create_futures_func = dask_create_futures_func(
+                client, run_func, name, retries=retries
+            )
+            create_backup_futures_func = dask_create_futures_func(
+                client, run_func, name
+            )
+
+            await async_map_dag(
+                create_futures_func,
+                dag=dag,
+                callbacks=callbacks,
+                resume=resume,
+                compute_arrays_in_parallel=compute_arrays_in_parallel,
+                create_backup_futures_func=create_backup_futures_func,
+                **kwargs,
+            )
diff --git a/cubed/runtime/executors/local.py b/cubed/runtime/executors/local.py
index a98c5fee..e68e0d04 100644
--- a/cubed/runtime/executors/local.py
+++ b/cubed/runtime/executors/local.py
@@ -1,26 +1,23 @@
 import asyncio
 import multiprocessing
 import os
-from concurrent.futures import Executor, ProcessPoolExecutor, ThreadPoolExecutor
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
 from functools import partial
-from typing import Any, AsyncIterator, Callable, Iterable, Optional, Sequence
+from typing import Any, Callable, Optional, Sequence
 
 import cloudpickle
 import psutil
-from aiostream import stream
-from aiostream.core import Stream
 from networkx import MultiDiGraph
 from tenacity import Retrying, stop_after_attempt
 
+from cubed.runtime.asyncio import async_map_dag
 from cubed.runtime.backup import use_backups_default
-from cubed.runtime.executors.asyncio import async_map_unordered
-from cubed.runtime.pipeline import visit_node_generations, visit_nodes
+from cubed.runtime.pipeline import visit_nodes
 from cubed.runtime.types import Callback, CubedPipeline, DagExecutor, TaskEndEvent
 from cubed.runtime.utils import (
     asyncio_run,
     execution_stats,
     execution_timing,
-    handle_callbacks,
     handle_operation_start_callbacks,
     profile_memray,
 )
@@ -83,92 +80,6 @@ def unpickle_and_call(f, inp, **kwargs):
     return f(inp, **kwargs)
 
 
-async def map_unordered(
-    concurrent_executor: Executor,
-    function: Callable[..., Any],
-    input: Iterable[Any],
-    retries: int = 2,
-    use_backups: bool = False,
-    batch_size: Optional[int] = None,
-    return_stats: bool = False,
-    name: Optional[str] = None,
-    **kwargs,
-) -> AsyncIterator[Any]:
-    if retries == 0:
-        retrying_function = function
-    else:
-        if isinstance(concurrent_executor, ProcessPoolExecutor):
-            raise NotImplementedError("Retries not supported for ProcessPoolExecutor")
-        retryer = Retrying(reraise=True, stop=stop_after_attempt(retries + 1))
-        retrying_function = partial(retryer, function)
-
-    def create_futures_func(input, **kwargs):
-        return [
-            (
-                i,
-                asyncio.wrap_future(
-                    concurrent_executor.submit(retrying_function, i, **kwargs)
-                ),
-            )
-            for i in input
-        ]
-
-    def create_futures_func_multiprocessing(input, **kwargs):
-        # Pickle the function, args, and kwargs using cloudpickle.
-        # They will be unpickled by unpickle_and_call.
-        pickled_kwargs = {k: cloudpickle.dumps(v) for k, v in kwargs.items()}
-        return [
-            (
-                i,
-                asyncio.wrap_future(
-                    concurrent_executor.submit(
-                        unpickle_and_call,
-                        cloudpickle.dumps(retrying_function),
-                        cloudpickle.dumps(i),
-                        **pickled_kwargs,
-                    )
-                ),
-            )
-            for i in input
-        ]
-
-    if isinstance(concurrent_executor, ProcessPoolExecutor):
-        create_futures = create_futures_func_multiprocessing
-    else:
-        create_futures = create_futures_func
-    async for result in async_map_unordered(
-        create_futures,
-        input,
-        use_backups=use_backups,
-        batch_size=batch_size,
-        return_stats=return_stats,
-        name=name,
-        **kwargs,
-    ):
-        yield result
-
-
-def pipeline_to_stream(
-    concurrent_executor: Executor,
-    run_func: Callable,
-    name: str,
-    pipeline: CubedPipeline,
-    **kwargs,
-) -> Stream:
-    return stream.iterate(
-        map_unordered(
-            concurrent_executor,
-            run_func,
-            pipeline.mappable,
-            return_stats=True,
-            name=name,
-            func=pipeline.function,
-            config=pipeline.config,
-            **kwargs,
-        )
-    )
-
-
 def check_runtime_memory(spec, max_workers):
     allowed_mem = spec.allowed_mem if spec is not None else None
     total_mem = psutil.virtual_memory().total
@@ -179,77 +90,30 @@ def check_runtime_memory(spec, max_workers):
             )
 
 
-async def async_execute_dag(
-    dag: MultiDiGraph,
-    callbacks: Optional[Sequence[Callback]] = None,
-    resume: Optional[bool] = None,
-    spec: Optional[Spec] = None,
-    compute_arrays_in_parallel: Optional[bool] = None,
-    **kwargs,
-) -> None:
-    concurrent_executor: Executor
-    max_workers = kwargs.pop("max_workers", os.cpu_count())
-    use_processes = kwargs.pop("use_processes", False)
-    if spec is not None:
-        check_runtime_memory(spec, max_workers)
-        if "use_backups" not in kwargs and use_backups_default(spec):
-            kwargs["use_backups"] = True
-    if use_processes:
-        max_tasks_per_child = kwargs.pop("max_tasks_per_child", None)
-        if isinstance(use_processes, str):
-            context = multiprocessing.get_context(use_processes)
-        else:
-            context = multiprocessing.get_context("spawn")
-        # max_tasks_per_child is only supported from Python 3.11
-        if max_tasks_per_child is None:
-            concurrent_executor = ProcessPoolExecutor(
-                max_workers=max_workers, mp_context=context
-            )
-        else:
-            concurrent_executor = ProcessPoolExecutor(
-                max_workers=max_workers,
-                mp_context=context,
-                max_tasks_per_child=max_tasks_per_child,
+def threads_create_futures_func(
+    concurrent_executor, function: Callable[..., Any], retries: int = 2
+):
+    if retries != 0:
+        retryer = Retrying(reraise=True, stop=stop_after_attempt(retries + 1))
+        function = partial(retryer, function)
+
+    def create_futures_func(input, **kwargs):
+        return [
+            (
+                i,
+                asyncio.wrap_future(concurrent_executor.submit(function, i, **kwargs)),
             )
-        run_func = run_func_processes
-    else:
-        concurrent_executor = ThreadPoolExecutor(max_workers=max_workers)
-        run_func = run_func_threads
-    try:
-        if not compute_arrays_in_parallel:
-            # run one pipeline at a time
-            for name, node in visit_nodes(dag, resume=resume):
-                handle_operation_start_callbacks(callbacks, name)
-                st = pipeline_to_stream(
-                    concurrent_executor, run_func, name, node["pipeline"], **kwargs
-                )
-                async with st.stream() as streamer:
-                    async for result, stats in streamer:
-                        handle_callbacks(callbacks, result, stats)
-        else:
-            for gen in visit_node_generations(dag, resume=resume):
-                # run pipelines in the same topological generation in parallel by merging their streams
-                streams = [
-                    pipeline_to_stream(
-                        concurrent_executor, run_func, name, node["pipeline"], **kwargs
-                    )
-                    for name, node in gen
-                ]
-                merged_stream = stream.merge(*streams)
-                async with merged_stream.stream() as streamer:
-                    async for result, stats in streamer:
-                        handle_callbacks(callbacks, result, stats)
+            for i in input
+        ]
 
-    finally:
-        # don't wait for any cancelled tasks
-        concurrent_executor.shutdown(wait=False)
+    return create_futures_func
 
 
 class ThreadsExecutor(DagExecutor):
     """An execution engine that uses Python asyncio."""
 
     def __init__(self, **kwargs):
-        self.kwargs = {**kwargs, **dict(use_processes=False)}
+        self.kwargs = kwargs
 
         # Tell NumPy to use a single thread
         # from https://stackoverflow.com/questions/30791550/limit-number-of-threads-in-numpy
@@ -273,7 +137,7 @@ class ThreadsExecutor(DagExecutor):
     ) -> None:
         merged_kwargs = {**self.kwargs, **kwargs}
         asyncio_run(
-            async_execute_dag(
+            self._async_execute_dag(
                 dag,
                 callbacks=callbacks,
                 resume=resume,
@@ -283,12 +147,67 @@ class ThreadsExecutor(DagExecutor):
             )
         )
 
+    async def _async_execute_dag(
+        self,
+        dag: MultiDiGraph,
+        callbacks: Optional[Sequence[Callback]] = None,
+        resume: Optional[bool] = None,
+        spec: Optional[Spec] = None,
+        compute_arrays_in_parallel: Optional[bool] = None,
+        **kwargs,
+    ) -> None:
+        max_workers = kwargs.pop("max_workers", os.cpu_count())
+        if spec is not None:
+            check_runtime_memory(spec, max_workers)
+            if "use_backups" not in kwargs and use_backups_default(spec):
+                kwargs["use_backups"] = True
+
+        concurrent_executor = ThreadPoolExecutor(max_workers=max_workers)
+        try:
+            create_futures_func = threads_create_futures_func(
+                concurrent_executor, run_func_threads, kwargs.pop("retries", 2)
+            )
+            await async_map_dag(
+                create_futures_func,
+                dag=dag,
+                callbacks=callbacks,
+                resume=resume,
+                compute_arrays_in_parallel=compute_arrays_in_parallel,
+                **kwargs,
+            )
+        finally:
+            # don't wait for any cancelled tasks
+            concurrent_executor.shutdown(wait=False)
+
+
+def processes_create_futures_func(concurrent_executor, function: Callable[..., Any]):
+    def create_futures_func(input, **kwargs):
+        # Pickle the function, args, and kwargs using cloudpickle.
+        # They will be unpickled by unpickle_and_call.
+        pickled_kwargs = {k: cloudpickle.dumps(v) for k, v in kwargs.items()}
+        return [
+            (
+                i,
+                asyncio.wrap_future(
+                    concurrent_executor.submit(
+                        unpickle_and_call,
+                        cloudpickle.dumps(function),
+                        cloudpickle.dumps(i),
+                        **pickled_kwargs,
+                    )
+                ),
+            )
+            for i in input
+        ]
+
+    return create_futures_func
+
 
 class ProcessesExecutor(DagExecutor):
     """An execution engine that uses local processes."""
 
     def __init__(self, **kwargs):
-        self.kwargs = {**kwargs, **dict(retries=0, use_processes=True)}
+        self.kwargs = kwargs
 
         # Tell NumPy to use a single thread
         # from https://stackoverflow.com/questions/30791550/limit-number-of-threads-in-numpy
@@ -312,7 +231,7 @@ class ProcessesExecutor(DagExecutor):
     ) -> None:
         merged_kwargs = {**self.kwargs, **kwargs}
         asyncio_run(
-            async_execute_dag(
+            self._async_execute_dag(
                 dag,
                 callbacks=callbacks,
                 resume=resume,
@@ -321,3 +240,52 @@ class ProcessesExecutor(DagExecutor):
                 **merged_kwargs,
             )
         )
+
+    async def _async_execute_dag(
+        self,
+        dag: MultiDiGraph,
+        callbacks: Optional[Sequence[Callback]] = None,
+        resume: Optional[bool] = None,
+        spec: Optional[Spec] = None,
+        compute_arrays_in_parallel: Optional[bool] = None,
+        **kwargs,
+    ) -> None:
+        max_workers = kwargs.pop("max_workers", os.cpu_count())
+        if spec is not None:
+            check_runtime_memory(spec, max_workers)
+            if "use_backups" not in kwargs and use_backups_default(spec):
+                kwargs["use_backups"] = True
+
+        use_processes = kwargs.pop("use_processes", True)
+        if isinstance(use_processes, str):
+            context = multiprocessing.get_context(use_processes)
+        else:
+            context = multiprocessing.get_context("spawn")
+
+        # max_tasks_per_child is only supported from Python 3.11
+        max_tasks_per_child = kwargs.pop("max_tasks_per_child", None)
+        if max_tasks_per_child is None:
+            concurrent_executor = ProcessPoolExecutor(
+                max_workers=max_workers, mp_context=context
+            )
+        else:
+            concurrent_executor = ProcessPoolExecutor(
+                max_workers=max_workers,
+                mp_context=context,
+                max_tasks_per_child=max_tasks_per_child,
+            )
+        try:
+            create_futures_func = processes_create_futures_func(
+                concurrent_executor, run_func_processes
+            )
+            await async_map_dag(
+                create_futures_func,
+                dag=dag,
+                callbacks=callbacks,
+                resume=resume,
+                compute_arrays_in_parallel=compute_arrays_in_parallel,
+                **kwargs,
+            )
+        finally:
+            # don't wait for any cancelled tasks
+            concurrent_executor.shutdown(wait=False)
diff --git a/cubed/runtime/executors/modal.py b/cubed/runtime/executors/modal.py
index 7b4a9676..b285445f 100644
--- a/cubed/runtime/executors/modal.py
+++ b/cubed/runtime/executors/modal.py
@@ -1,26 +1,17 @@
 import asyncio
 import os
-import time
 from asyncio.exceptions import TimeoutError
-from typing import Any, AsyncIterator, Iterable, Optional, Sequence
+from typing import Any, Callable, Optional, Sequence
 
 import modal
-from aiostream import stream
 from modal.exception import ConnectionError
-from modal.functions import Function
 from networkx import MultiDiGraph
 from tenacity import retry, retry_if_exception_type, stop_after_attempt
 
+from cubed.runtime.asyncio import async_map_dag
 from cubed.runtime.backup import use_backups_default
-from cubed.runtime.executors.asyncio import async_map_unordered
-from cubed.runtime.pipeline import visit_node_generations, visit_nodes
 from cubed.runtime.types import Callback, DagExecutor
-from cubed.runtime.utils import (
-    asyncio_run,
-    execute_with_stats,
-    handle_callbacks,
-    handle_operation_start_callbacks,
-)
+from cubed.runtime.utils import asyncio_run, execute_with_stats
 from cubed.spec import Spec
 
 RUNTIME_MEMORY_MIB = 2000
@@ -117,129 +108,13 @@ class Container:
         return result, stats
 
 
-# We need map_unordered for the use_backups implementation
-async def map_unordered(
-    app_function: Function,
-    input: Iterable[Any],
-    use_backups: bool = False,
-    backup_function: Optional[Function] = None,
-    batch_size: Optional[int] = None,
-    return_stats: bool = False,
-    name: Optional[str] = None,
-    **kwargs,
-) -> AsyncIterator[Any]:
-    """
-    Apply a function to items of an input list, yielding results as they are completed
-    (which may be different to the input order).
-
-    :param app_function: The Modal function to map over the data.
-    :param input: An iterable of input data.
-    :param use_backups: Whether to launch backup tasks to mitigate against slow-running tasks.
-    :param kwargs: Keyword arguments to pass to the function.
-
-    :return: Function values (and optionally stats) as they are completed, not necessarily in the input order.
-    """
-
-    if not use_backups and batch_size is None:
-        task_create_tstamp = time.time()
-        async for result in app_function.map(input, order_outputs=False, kwargs=kwargs):
-            if return_stats:
-                result, stats = result
-                if name is not None:
-                    stats["name"] = name
-                stats["task_create_tstamp"] = task_create_tstamp
-                yield result, stats
-            else:
-                yield result
-        return
-
+def modal_create_futures_func(function: Callable[..., Any]):
     def create_futures_func(input, **kwargs):
         return [
-            (i, asyncio.ensure_future(app_function.remote.aio(i, **kwargs)))
-            for i in input
+            (i, asyncio.ensure_future(function.remote.aio(i, **kwargs))) for i in input
         ]
 
-    backup_function = backup_function or app_function
-
-    def create_backup_futures_func(input, **kwargs):
-        return [
-            (i, asyncio.ensure_future(backup_function.remote.aio(i, **kwargs)))
-            for i in input
-        ]
-
-    async for result in async_map_unordered(
-        create_futures_func,
-        input,
-        use_backups=use_backups,
-        create_backup_futures_func=create_backup_futures_func,
-        batch_size=batch_size,
-        return_stats=return_stats,
-        name=name,
-        **kwargs,
-    ):
-        yield result
-
-
-def pipeline_to_stream(app_function, name, pipeline, **kwargs):
-    return stream.iterate(
-        map_unordered(
-            app_function,
-            pipeline.mappable,
-            return_stats=True,
-            name=name,
-            func=pipeline.function,
-            config=pipeline.config,
-            **kwargs,
-        )
-    )
-
-
-# This just retries the initial connection attempt, not the function calls
-@retry(
-    reraise=True,
-    retry=retry_if_exception_type((TimeoutError, ConnectionError)),
-    stop=stop_after_attempt(3),
-)
-async def async_execute_dag(
-    dag: MultiDiGraph,
-    callbacks: Optional[Sequence[Callback]] = None,
-    resume: Optional[bool] = None,
-    spec: Optional[Spec] = None,
-    cloud: Optional[str] = None,
-    compute_arrays_in_parallel: Optional[bool] = None,
-    **kwargs,
-) -> None:
-    if spec is not None:
-        check_runtime_memory(spec)
-        if "use_backups" not in kwargs and use_backups_default(spec):
-            kwargs["use_backups"] = True
-    async with app.run():
-        cloud = cloud or "aws"
-        if cloud == "aws":
-            app_function = run_remotely
-        elif cloud == "gcp":
-            app_function = Container().run_remotely
-        else:
-            raise ValueError(f"Unrecognized cloud: {cloud}")
-        if not compute_arrays_in_parallel:
-            # run one pipeline at a time
-            for name, node in visit_nodes(dag, resume=resume):
-                handle_operation_start_callbacks(callbacks, name)
-                st = pipeline_to_stream(app_function, name, node["pipeline"], **kwargs)
-                async with st.stream() as streamer:
-                    async for result, stats in streamer:
-                        handle_callbacks(callbacks, result, stats)
-        else:
-            for gen in visit_node_generations(dag, resume=resume):
-                # run pipelines in the same topological generation in parallel by merging their streams
-                streams = [
-                    pipeline_to_stream(app_function, name, node["pipeline"], **kwargs)
-                    for name, node in gen
-                ]
-                merged_stream = stream.merge(*streams)
-                async with merged_stream.stream() as streamer:
-                    async for result, stats in streamer:
-                        handle_callbacks(callbacks, result, stats)
+    return create_futures_func
 
 
 class ModalExecutor(DagExecutor):
@@ -263,7 +138,7 @@ class ModalExecutor(DagExecutor):
     ) -> None:
         merged_kwargs = {**self.kwargs, **kwargs}
         asyncio_run(
-            async_execute_dag(
+            self._async_execute_dag(
                 dag,
                 callbacks=callbacks,
                 resume=resume,
@@ -272,3 +147,43 @@ class ModalExecutor(DagExecutor):
                 **merged_kwargs,
             )
         )
+
+    # This just retries the initial connection attempt, not the function calls
+    @retry(
+        reraise=True,
+        retry=retry_if_exception_type((TimeoutError, ConnectionError)),
+        stop=stop_after_attempt(3),
+    )
+    async def _async_execute_dag(
+        self,
+        dag: MultiDiGraph,
+        callbacks: Optional[Sequence[Callback]] = None,
+        resume: Optional[bool] = None,
+        spec: Optional[Spec] = None,
+        cloud: Optional[str] = None,
+        compute_arrays_in_parallel: Optional[bool] = None,
+        **kwargs,
+    ) -> None:
+        if spec is not None:
+            check_runtime_memory(spec)
+            if "use_backups" not in kwargs and use_backups_default(spec):
+                kwargs["use_backups"] = True
+
+        cloud = cloud or "aws"
+        if cloud == "aws":
+            function = run_remotely
+        elif cloud == "gcp":
+            function = Container().run_remotely
+        else:
+            raise ValueError(f"Unrecognized cloud: {cloud}")
+
+        async with app.run():
+            create_futures_func = modal_create_futures_func(function)
+            await async_map_dag(
+                create_futures_func,
+                dag=dag,
+                callbacks=callbacks,
+                resume=resume,
+                compute_arrays_in_parallel=compute_arrays_in_parallel,
+                **kwargs,
+            )


Test patch:
diff --git a/cubed/tests/runtime/test_dask.py b/cubed/tests/runtime/test_dask.py
index 13eb2e49..c2610fa7 100644
--- a/cubed/tests/runtime/test_dask.py
+++ b/cubed/tests/runtime/test_dask.py
@@ -4,23 +4,25 @@ from functools import partial
 
 import pytest
 
+from cubed.runtime.asyncio import async_map_unordered
 from cubed.tests.runtime.utils import check_invocation_counts, deterministic_failure
 
 pytest.importorskip("dask.distributed")
 
 from dask.distributed import Client
 
-from cubed.runtime.executors.dask import map_unordered
+from cubed.runtime.executors.dask import dask_create_futures_func
 
 
 async def run_test(function, input, retries, use_backups=False, batch_size=None):
     outputs = set()
     async with Client(asynchronous=True) as client:
-        async for output in map_unordered(
-            client,
-            function,
+        create_futures_func = dask_create_futures_func(
+            client, function, retries=retries
+        )
+        async for output in async_map_unordered(
+            create_futures_func,
             input,
-            retries=retries,
             use_backups=use_backups,
             batch_size=batch_size,
         ):
diff --git a/cubed/tests/runtime/test_local.py b/cubed/tests/runtime/test_local.py
index 99e3fad1..a1a0a9f0 100644
--- a/cubed/tests/runtime/test_local.py
+++ b/cubed/tests/runtime/test_local.py
@@ -5,7 +5,8 @@ from functools import partial
 
 import pytest
 
-from cubed.runtime.executors.local import map_unordered
+from cubed.runtime.asyncio import async_map_unordered
+from cubed.runtime.executors.local import threads_create_futures_func
 from cubed.tests.runtime.utils import check_invocation_counts, deterministic_failure
 
 
@@ -13,11 +14,12 @@ async def run_test(function, input, retries=2, use_backups=False, batch_size=Non
     outputs = set()
     concurrent_executor = ThreadPoolExecutor()
     try:
-        async for output in map_unordered(
-            concurrent_executor,
-            function,
+        create_futures_func = threads_create_futures_func(
+            concurrent_executor, function, retries=retries
+        )
+        async for output in async_map_unordered(
+            create_futures_func,
             input,
-            retries=retries,
             use_backups=use_backups,
             batch_size=batch_size,
         ):
diff --git a/cubed/tests/runtime/test_modal.py b/cubed/tests/runtime/test_modal.py
index 6e341709..a92dbaeb 100644
--- a/cubed/tests/runtime/test_modal.py
+++ b/cubed/tests/runtime/test_modal.py
@@ -9,7 +9,8 @@ import asyncio
 import fsspec
 import modal
 
-from cubed.runtime.executors.modal import map_unordered
+from cubed.runtime.asyncio import async_map_unordered
+from cubed.runtime.executors.modal import modal_create_futures_func
 from cubed.tests.runtime.utils import check_invocation_counts, deterministic_failure
 
 tmp_path = "s3://cubed-unittest/map_unordered"
@@ -66,8 +67,9 @@ def deterministic_failure_modal_long_timeout(
 async def run_test(app_function, input, use_backups=False, batch_size=None, **kwargs):
     outputs = set()
     async with app.run():
-        async for output in map_unordered(
-            app_function,
+        create_futures_func = modal_create_futures_func(app_function)
+        async for output in async_map_unordered(
+            create_futures_func,
             input,
             use_backups=use_backups,
             batch_size=batch_size,
diff --git a/cubed/tests/test_array_api.py b/cubed/tests/test_array_api.py
index a0b023b2..1c110e64 100644
--- a/cubed/tests/test_array_api.py
+++ b/cubed/tests/test_array_api.py
@@ -194,6 +194,19 @@ def test_add_different_chunks_fail(spec, executor):
     assert_array_equal(c.compute(executor=executor), np.ones((10,)) + np.ones((10,)))
 
 
+def test_add_scalars():
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2))
+
+    b = xp.add(a, 1)
+    assert_array_equal(b.compute(), np.array([[2, 3, 4], [5, 6, 7], [8, 9, 10]]))
+
+    c = xp.add(2, a)
+    assert_array_equal(c.compute(), np.array([[3, 4, 5], [6, 7, 8], [9, 10, 11]]))
+
+    with pytest.raises(TypeError):
+        xp.add(1, 2)
+
+
 @pytest.mark.parametrize(
     "min, max",
     [
diff --git a/cubed/tests/utils.py b/cubed/tests/utils.py
index 6e66e585..e3d266bb 100644
--- a/cubed/tests/utils.py
+++ b/cubed/tests/utils.py
@@ -27,6 +27,7 @@ MAIN_EXECUTORS = [create_executor("single-threaded")]
 if platform.system() != "Windows":
     # ThreadsExecutor calls `peak_measured_mem` which is not supported on Windows
     ALL_EXECUTORS.append(create_executor("threads"))
+    MAIN_EXECUTORS.append(create_executor("threads"))
 
     ALL_EXECUTORS.append(create_executor("processes"))
     MAIN_EXECUTORS.append(create_executor("processes"))