Repository: cubed-dev/cubed
Pull request #595 (2024-10-14 16:28:34)

Add `map_selection` primitive to replace `map_direct`
This is still a WIP, but the basic idea is to migrate calls to `map_direct` to `map_selection`. The latter maps output block keys to an index selection on the input arrays, and is implemented using blockwise, so the optimizer will be able to fuse operations together. See #464.


Patch:
diff --git a/cubed/array_api/manipulation_functions.py b/cubed/array_api/manipulation_functions.py
index a18a133a..6229f624 100644
--- a/cubed/array_api/manipulation_functions.py
+++ b/cubed/array_api/manipulation_functions.py
@@ -9,7 +9,13 @@ from cubed.array_api.creation_functions import empty
 from cubed.backend_array_api import namespace as nxp
 from cubed.core import squeeze  # noqa: F401
 from cubed.core import blockwise, rechunk, unify_chunks
-from cubed.core.ops import elemwise, general_blockwise, map_blocks, map_direct
+from cubed.core.ops import (
+    elemwise,
+    general_blockwise,
+    map_blocks,
+    map_direct,
+    map_selection,
+)
 from cubed.utils import block_id_to_offset, get_item, offset_to_block_id, to_chunksize
 from cubed.vendor.dask.array.core import broadcast_chunks, normalize_chunks
 from cubed.vendor.dask.array.reshape import reshape_rechunk
@@ -178,42 +184,50 @@ def flip(x, /, *, axis=None):
     if not isinstance(axis, tuple):
         axis = (axis,)
     axis = validate_axis(axis, x.ndim)
-    return map_direct(
-        _flip,
+
+    def selection_function(out_key):
+        out_coords = out_key[1:]
+        block_id = out_coords
+
+        # produce a key that has slices (except for axis dimensions, which are replaced below)
+        idx = tuple(0 if i == axis else v for i, v in enumerate(block_id))
+        key = list(get_item(x.chunks, idx))
+
+        for ax in axis:
+            # determine the start and stop indexes for this block along the axis dimension
+            start = block_id[ax] * x.chunksize[ax]
+            stop = start + x.chunksize[ax]
+            stop = min(stop, x.shape[ax])
+
+            # flip start and stop
+            axis_len = x.shape[ax]
+            start, stop = axis_len - stop, axis_len - start
+
+            # replace with slice
+            key[ax] = slice(start, stop)
+
+        return tuple(key)
+
+    max_num_input_blocks = _flip_num_input_blocks(axis, x.shape, x.chunksize)
+
+    return map_selection(
+        nxp.flip,
+        selection_function,
         x,
         shape=x.shape,
         dtype=x.dtype,
         chunks=x.chunks,
-        extra_projected_mem=x.chunkmem,
-        target_chunks=x.chunks,
+        max_num_input_blocks=max_num_input_blocks,
         axis=axis,
     )
 
 
-def _flip(x, *arrays, target_chunks=None, axis=None, block_id=None):
-    array = arrays[0].zarray  # underlying Zarr array (or virtual array)
-    chunks = target_chunks
-
-    # produce a key that has slices (except for axis dimensions, which are replaced below)
-    idx = tuple(0 if i == axis else v for i, v in enumerate(block_id))
-    key = list(get_item(chunks, idx))
-
+def _flip_num_input_blocks(axis, shape, chunksizes):
+    num = 1
     for ax in axis:
-        # determine the start and stop indexes for this block along the axis dimension
-        chunksize = to_chunksize(chunks)
-        start = block_id[ax] * chunksize[ax]
-        stop = start + x.shape[ax]
-
-        # flip start and stop
-        axis_len = array.shape[ax]
-        start, stop = axis_len - stop, axis_len - start
-
-        # replace with slice
-        key[ax] = slice(start, stop)
-
-    key = tuple(key)
-
-    return nxp.flip(array[key], axis=axis)
+        if shape[ax] % chunksizes[ax] != 0:
+            num *= 2
+    return num
 
 
 def moveaxis(
diff --git a/cubed/core/groupby.py b/cubed/core/groupby.py
index 8c4aa9b7..8ec97045 100644
--- a/cubed/core/groupby.py
+++ b/cubed/core/groupby.py
@@ -2,8 +2,8 @@ from typing import TYPE_CHECKING
 
 from cubed.array_api.manipulation_functions import broadcast_to, expand_dims
 from cubed.backend_array_api import namespace as nxp
-from cubed.core.ops import map_blocks, map_direct, reduction
-from cubed.utils import array_memory, get_item
+from cubed.core.ops import map_blocks, map_selection, reduction
+from cubed.utils import get_item
 from cubed.vendor.dask.array.core import normalize_chunks
 
 if TYPE_CHECKING:
@@ -181,57 +181,51 @@ def groupby_blockwise(
     )
     target_chunks = normalize_chunks(chunks, shape, dtype=dtype)
 
-    # memory allocated by reading one chunk from input array
-    # note that although read_chunks will overlap multiple input chunks, zarr will
-    # read the chunks in series, reusing the buffer
-    extra_projected_mem = x.chunkmem
+    def selection_function(out_key):
+        out_coords = out_key[1:]
+        block_id = out_coords
+        return get_item(read_chunks, block_id)
 
-    # memory allocated for largest of (variable sized) read_chunks
-    read_chunksize = tuple(max(c) for c in read_chunks)
-    extra_projected_mem += array_memory(x.dtype, read_chunksize)
+    # in general each selection overlaps 2 input blocks
+    max_num_input_blocks = 2
 
-    return map_direct(
+    return map_selection(
         _process_blockwise_chunk,
+        selection_function,
         x,
         shape=shape,
         dtype=dtype,
         chunks=target_chunks,
-        extra_projected_mem=extra_projected_mem,
+        max_num_input_blocks=max_num_input_blocks,
         axis=axis,
         by=by,
         blockwise_func=func,
-        read_chunks=read_chunks,
         by_read_chunks=by_read_chunks,
         target_chunks=target_chunks,
         groups_per_chunk=groups_per_chunk,
-        extra_func_kwargs=extra_func_kwargs,
     )
 
 
 def _process_blockwise_chunk(
-    x,
-    *arrays,
+    a,
     axis=None,
     by=None,
     blockwise_func=None,
-    read_chunks=None,
     by_read_chunks=None,
     target_chunks=None,
     groups_per_chunk=None,
     block_id=None,
     **kwargs,
 ):
-    array = arrays[0].zarray  # underlying Zarr array (or virtual array)
     idx = block_id
     bi = idx[axis]
 
-    result = array[get_item(read_chunks, idx)]
     by = by[get_item(by_read_chunks, (bi,))]
 
     start_group = bi * groups_per_chunk
 
     return blockwise_func(
-        result,
+        a,
         by,
         axis=axis,
         start_group=start_group,
diff --git a/cubed/core/ops.py b/cubed/core/ops.py
index d82c4d03..39121fb2 100644
--- a/cubed/core/ops.py
+++ b/cubed/core/ops.py
@@ -286,6 +286,7 @@ def blockwise(
 
     fusable = kwargs.pop("fusable", True)
     num_input_blocks = kwargs.pop("num_input_blocks", None)
+    iterable_input_blocks = kwargs.pop("iterable_input_blocks", None)
 
     name = gensym()
     spec = check_array_specs(arrays)
@@ -311,6 +312,7 @@ def blockwise(
         extra_func_kwargs=extra_func_kwargs,
         fusable=fusable,
         num_input_blocks=num_input_blocks,
+        iterable_input_blocks=iterable_input_blocks,
         **kwargs,
     )
     plan = Plan._new(
@@ -368,6 +370,11 @@ def general_blockwise(
         num_input_blocks = kwargs.pop("num_input_blocks", None)
         if num_input_blocks is not None:
             num_input_blocks = num_input_blocks + (1,)  # for offsets array
+        iterable_input_blocks = kwargs.pop("iterable_input_blocks", None)
+        if iterable_input_blocks is not None:
+            iterable_input_blocks = iterable_input_blocks + (
+                False,
+            )  # for offsets array
 
         return _general_blockwise(
             func_with_block_id(func),
@@ -380,6 +387,7 @@ def general_blockwise(
             target_paths=target_paths,
             extra_func_kwargs=extra_func_kwargs,
             num_input_blocks=num_input_blocks,
+            iterable_input_blocks=iterable_input_blocks,
             **kwargs,
         )
 
@@ -421,6 +429,7 @@ def _general_blockwise(
     extra_projected_mem = kwargs.pop("extra_projected_mem", 0)
 
     num_input_blocks = kwargs.pop("num_input_blocks", None)
+    iterable_input_blocks = kwargs.pop("iterable_input_blocks", None)
 
     spec = check_array_specs(arrays)
 
@@ -452,6 +461,7 @@ def _general_blockwise(
         in_names=in_names,
         extra_func_kwargs=extra_func_kwargs,
         num_input_blocks=num_input_blocks,
+        iterable_input_blocks=iterable_input_blocks,
         **kwargs,
     )
     plan = Plan._new(
@@ -565,37 +575,33 @@ def index(x, key):
             if not isinstance(ia, ndindex.Integer)
         )
 
+        # this is the same as chunks, except it has the same number of dimensions as the input
+        out_chunksizes = tuple(
+            chunk_len_for_indexer(ia, c) if not isinstance(ia, ndindex.Integer) else 1
+            for ia, c in zip(idx.args, x.chunksize)
+        )
+
         target_chunks = normalize_chunks(chunks, shape, dtype=dtype)
 
-        if _is_chunk_aligned_selection(idx):
-            # use general_blockwise, which allows more opportunities for optimization than map_direct
+        if _is_basic_selection(idx):
+            # use map_selection (which uses general_blockwise) to allow more opportunities for optimization than map_direct
 
-            def key_function(out_key):
+            def selection_function(out_key):
                 out_coords = out_key[1:]
+                return _target_chunk_selection(target_chunks, out_coords, selection)
 
-                # compute the selection on x required to get the relevant chunk for out_coords
-                in_sel = _target_chunk_selection(target_chunks, out_coords, selection)
-
-                # use a Zarr BasicIndexer to convert this to input coordinates
-                indexer = create_basic_indexer(
-                    in_sel, x.zarray_maybe_lazy.shape, x.zarray_maybe_lazy.chunks
-                )
-
-                return tuple(
-                    (x.name,) + chunk_coords for (chunk_coords, _, _) in indexer
-                )
+            max_num_input_blocks = _index_num_input_blocks(
+                idx, x.chunksize, out_chunksizes, x.numblocks
+            )
 
-            out = general_blockwise(
-                _assemble_index_chunk,
-                key_function,
+            out = map_selection(
+                None,  # no function to apply after selection
+                selection_function,
                 x,
-                shapes=[shape],
-                dtypes=[x.dtype],
-                chunkss=[target_chunks],
-                target_chunks=target_chunks,
-                selection=selection,
-                in_shape=x.shape,
-                in_chunksize=x.chunksize,
+                shape,
+                x.dtype,
+                target_chunks,
+                max_num_input_blocks=max_num_input_blocks,
             )
         else:
             # use map_direct, which can't be fused
@@ -635,16 +641,29 @@ def index(x, key):
     return out
 
 
-def _is_chunk_aligned_selection(idx: ndindex.Tuple):
-    return all(
-        isinstance(ia, ndindex.Integer)
-        or (
-            isinstance(ia, ndindex.Slice)
-            and ia.start == 0
-            and (ia.step is None or ia.step == 1)
-        )
-        for ia in idx.args
-    )
+def _is_basic_selection(idx: ndindex.Tuple):
+    return all(isinstance(ia, (ndindex.Integer, ndindex.Slice)) for ia in idx.args)
+
+
+def _index_num_input_blocks(
+    idx: ndindex.Tuple, in_chunksizes, out_chunksizes, numblocks
+):
+    num = 1
+    for ia, c, oc, nb in zip(idx.args, in_chunksizes, out_chunksizes, numblocks):
+        if isinstance(ia, ndindex.Integer) or nb == 1:
+            pass  # single block
+        elif isinstance(ia, ndindex.Slice):
+            if (ia.start // c) == ((ia.stop - 1) // c):
+                pass  # within same block
+            elif ia.start % c != 0:
+                num *= 2  # doesn't start on chunk boundary
+            elif ia.step is not None and c % ia.step != 0 and oc > 1:
+                # step is not a multiple of chunk size, and output chunks have more than one element
+                # so some output chunks will access two input chunks
+                num *= 2
+        else:
+            raise NotImplementedError("Only integer or slice indexes are supported.")
+    return num
 
 
 def create_basic_indexer(selection, shape, chunks):
@@ -670,28 +689,41 @@ class ZarrArrayIndexingAdaptor:
 
 
 def _assemble_index_chunk(
-    *arrs,
-    target_chunks=None,
-    selection=None,
+    arrays,
+    dtype=None,
+    func=None,
+    selection_function=None,
     in_shape=None,
     in_chunksize=None,
     block_id=None,
+    **kwargs,
 ):
+    assert not isinstance(
+        arrays, list
+    ), "index expects an iterator of array blocks, not a list"
+
     # compute the selection on x required to get the relevant chunk for out_coords
     out_coords = block_id
-    in_sel = _target_chunk_selection(target_chunks, out_coords, selection)
+    in_sel = selection_function(("out",) + out_coords)
 
     # use a Zarr BasicIndexer to convert this to input coordinates
     indexer = create_basic_indexer(in_sel, in_shape, in_chunksize)
 
     shape = indexer.shape
-    out = np.empty_like(arrs[0], shape=shape)
+    out = np.empty(shape, dtype=dtype)
 
     if array_size(shape) > 0:
         _, lchunk_selection, lout_selection = zip(*indexer)
-        for ai, chunk_select, out_select in zip(arrs, lchunk_selection, lout_selection):
+        for ai, chunk_select, out_select in zip(
+            arrays, lchunk_selection, lout_selection
+        ):
             out[out_select] = ai[chunk_select]
 
+    if func is not None:
+        if has_keyword(func, "block_id"):
+            out = func(out, block_id=block_id, **kwargs)
+        else:
+            out = func(out, **kwargs)
     return out
 
 
@@ -747,6 +779,51 @@ def _target_chunk_selection(target_chunks, idx, selection):
     return tuple(sel)
 
 
+def map_selection(
+    func,
+    selection_function,
+    x,
+    shape,
+    dtype,
+    chunks,
+    max_num_input_blocks,
+    **kwargs,
+) -> "Array":
+    def key_function(out_key):
+        # compute the selection on x required to get the relevant chunk for out_key
+        in_sel = selection_function(out_key)
+
+        # use a Zarr BasicIndexer to convert selection to input coordinates
+        indexer = create_basic_indexer(in_sel, x.shape, x.chunksize)
+
+        return (
+            iter(tuple((x.name,) + chunk_coords for (chunk_coords, _, _) in indexer)),
+        )
+
+    num_input_blocks = (max_num_input_blocks,)
+    iterable_input_blocks = (True,)
+
+    out = general_blockwise(
+        _assemble_index_chunk,
+        key_function,
+        x,
+        shapes=[shape],
+        dtypes=[x.dtype],
+        chunkss=[chunks],
+        extra_func_kwargs=dict(func=func, dtype=dtype),
+        num_input_blocks=num_input_blocks,
+        iterable_input_blocks=iterable_input_blocks,
+        selection_function=selection_function,
+        in_shape=x.shape,
+        in_chunksize=x.chunksize,
+        **kwargs,
+    )
+    from cubed import Array
+
+    assert isinstance(out, Array)  # single output
+    return out
+
+
 def map_blocks(
     func,
     *args: "Array",
@@ -1085,9 +1162,10 @@ def merge_chunks_new(x, chunks):
         # return a tuple with a single item that is the list of input keys to be merged
         return (lol_product((x.name,), in_keys),)
 
-    num_input_blocks = int(
-        np.prod([c1 // c0 for (c0, c1) in zip(x.chunksize, target_chunksize)])
+    num_input_blocks = (
+        int(np.prod([c1 // c0 for (c0, c1) in zip(x.chunksize, target_chunksize)])),
     )
+    iterable_input_blocks = (True,)
 
     return general_blockwise(
         _concatenate2,
@@ -1097,7 +1175,8 @@ def merge_chunks_new(x, chunks):
         dtypes=[x.dtype],
         chunkss=[target_chunks],
         extra_projected_mem=0,
-        num_input_blocks=(num_input_blocks,),
+        num_input_blocks=num_input_blocks,
+        iterable_input_blocks=iterable_input_blocks,
         axes=axes,
     )
 
@@ -1325,6 +1404,7 @@ def partial_reduce(
         chunkss=[chunks],
         extra_projected_mem=extra_projected_mem,
         num_input_blocks=(sum(split_every.values()),),
+        iterable_input_blocks=(True,),
         reduce_func=func,
         initial_func=initial_func,
         axis=axis,
diff --git a/cubed/overlap.py b/cubed/overlap.py
index 8ae2785c..1737f1d4 100644
--- a/cubed/overlap.py
+++ b/cubed/overlap.py
@@ -1,7 +1,7 @@
 from typing import Tuple
 
 from cubed.backend_array_api import namespace as nxp
-from cubed.core.ops import map_direct
+from cubed.core.ops import map_selection
 from cubed.types import T_RectangularChunks
 from cubed.utils import _cumsum
 from cubed.vendor.dask.array.core import normalize_chunks
@@ -57,47 +57,57 @@ def map_overlap(
     depth = coerce(args, depth, coerce_depth)
     boundary = coerce(args, boundary, coerce_boundary)
 
-    # memory allocated by reading one chunk from input array
-    # note that although the output chunk will overlap multiple input chunks, zarr will
-    # read the chunks in series, reusing the buffer
-    extra_projected_mem = args[0].chunkmem  # TODO: support multiple
+    x = args[0]  # TODO: support multiple input arrays
+
+    def selection_function(out_key):
+        out_coords = out_key[1:]
+        block_id = out_coords
+        return get_item_with_depth(x.chunks, block_id, depth[0])
+
+    max_num_input_blocks = _overlap_num_input_blocks(depth[0], x.numblocks)
 
     has_block_id_kw = has_keyword(func, "block_id")
 
-    return map_direct(
+    # First read the chunk with overlaps determined by depth, then pad boundaries second.
+    # Do it this way round so we can do everything with one blockwise. The alternative,
+    # which pads the entire array first (via concatenate), would result in at least one extra copy.
+
+    return map_selection(
         _overlap,
-        *args,
+        selection_function,
+        x,
         shape=shape,
         dtype=dtype,
         chunks=chunks,
-        extra_projected_mem=extra_projected_mem,
+        max_num_input_blocks=max_num_input_blocks,
         overlap_func=func,
         depth=depth,
         boundary=boundary,
+        numblocks=x.numblocks,
         has_block_id_kw=has_block_id_kw,
-        **kwargs,
     )
 
 
+def _overlap_num_input_blocks(depth, numblocks):
+    num = 1
+    for i in depth.keys():
+        num *= min(numblocks[i], 3)
+    return num
+
+
 def _overlap(
-    x,
-    *arrays,
+    a,
     overlap_func=None,
     depth=None,
     boundary=None,
+    numblocks=None,
     has_block_id_kw=False,
     block_id=None,
     **kwargs,
 ):
-    a = arrays[0]  # TODO: support multiple
     depth = depth[0]
     boundary = boundary[0]
-
-    # First read the chunk with overlaps determined by depth, then pad boundaries second.
-    # Do it this way round so we can do everything with one blockwise. The alternative,
-    # which pads the entire array first (via concatenate), would result in at least one extra copy.
-    out = a.zarray[get_item_with_depth(a.chunks, block_id, depth)]
-    out = _pad_boundaries(out, depth, boundary, a.numblocks, block_id)
+    out = _pad_boundaries(a, depth, boundary, numblocks, block_id)
     if has_block_id_kw:
         return overlap_func(out, block_id=block_id, **kwargs)
     else:
diff --git a/cubed/primitive/blockwise.py b/cubed/primitive/blockwise.py
index 1d6618c2..9c3b8495 100644
--- a/cubed/primitive/blockwise.py
+++ b/cubed/primitive/blockwise.py
@@ -55,6 +55,8 @@ class BlockwiseSpec:
         The number of array arguments that ``function`` takes.
     num_input_blocks: Tuple[int, ...]
         The number of input blocks read from each input array.
+    iterable_input_blocks: Tuple[int, ...]
+        Whether the input blocks read from each input array are supplied as an iterable or not.
     reads_map : Dict[str, CubedArrayProxy]
         Read proxy dictionary keyed by array name.
     writes_list : List[CubedArrayProxy]
@@ -65,6 +67,7 @@ class BlockwiseSpec:
     function: Callable[..., Any]
     function_nargs: int
     num_input_blocks: Tuple[int, ...]
+    iterable_input_blocks: Tuple[bool, ...]
     reads_map: Dict[str, CubedArrayProxy]
     writes_list: List[CubedArrayProxy]
 
@@ -143,6 +146,7 @@ def blockwise(
     extra_func_kwargs: Optional[Dict[str, Any]] = None,
     fusable: bool = True,
     num_input_blocks: Optional[Tuple[int, ...]] = None,
+    iterable_input_blocks: Optional[Tuple[bool, ...]] = None,
     **kwargs,
 ) -> PrimitiveOperation:
     """Apply a function to multiple blocks from multiple inputs, expressed using concise indexing rules.
@@ -229,6 +233,7 @@ def blockwise(
         extra_func_kwargs=extra_func_kwargs,
         fusable=fusable,
         num_input_blocks=num_input_blocks,
+        iterable_input_blocks=iterable_input_blocks,
         **kwargs,
     )
 
@@ -251,6 +256,7 @@ def general_blockwise(
     extra_func_kwargs: Optional[Dict[str, Any]] = None,
     fusable: bool = True,
     num_input_blocks: Optional[Tuple[int, ...]] = None,
+    iterable_input_blocks: Optional[Tuple[bool, ...]] = None,
     **kwargs,
 ) -> PrimitiveOperation:
     """A more general form of ``blockwise`` that uses a function to specify the block
@@ -297,6 +303,7 @@ def general_blockwise(
     func_kwargs = extra_func_kwargs or {}
     func_with_kwargs = partial(func, **{**kwargs, **func_kwargs})
     num_input_blocks = num_input_blocks or (1,) * len(arrays)
+    iterable_input_blocks = iterable_input_blocks or (False,) * len(arrays)
     read_proxies = {
         name: CubedArrayProxy(array, array.chunks) for name, array in array_map.items()
     }
@@ -346,6 +353,7 @@ def general_blockwise(
         func_with_kwargs,
         len(arrays),
         num_input_blocks,
+        iterable_input_blocks,
         read_proxies,
         write_proxies,
     )
@@ -433,18 +441,6 @@ def can_fuse_multiple_primitive_ops(
                 primitive_op.allowed_mem,
             )
             return False
-        # If the number of input blocks for each input is not uniform, then we
-        # can't fuse. (This should never happen since all operations are
-        # currently uniform, and fused operations are too if fuse is applied in
-        # topological order.)
-        num_input_blocks = primitive_op.pipeline.config.num_input_blocks
-        if not all(num_input_blocks[0] == n for n in num_input_blocks):
-            logger.debug(
-                "can't fuse %s since number of input blocks for each input is not uniform: %s",
-                name,
-                num_input_blocks,
-            )
-            return False
         if max_total_num_input_blocks is None:
             # If max total input blocks not specified, then only fuse if num
             # tasks of predecessor ops match.
@@ -462,6 +458,7 @@ def can_fuse_multiple_primitive_ops(
                 )
             return ret
         else:
+            num_input_blocks = primitive_op.pipeline.config.num_input_blocks
             total_num_input_blocks = 0
             for ni, p in zip(num_input_blocks, predecessor_primitive_ops):
                 for nj in p.pipeline.config.num_input_blocks:
@@ -526,11 +523,13 @@ def fuse(
         n * pipeline2.config.num_input_blocks[0]
         for n in pipeline1.config.num_input_blocks
     )
+    iterable_input_blocks = pipeline1.config.iterable_input_blocks
     spec = BlockwiseSpec(
         fused_key_func,
         fused_func,
         function_nargs,
         num_input_blocks,
+        iterable_input_blocks,
         read_proxies,
         write_proxies,
     )
@@ -567,113 +566,25 @@ def fuse_multiple(
     Fuse a blockwise operation and its predecessors into a single operation, avoiding writing to (or reading from) the targets of the predecessor operations.
     """
 
-    pipeline = primitive_op.pipeline
-    predecessor_pipelines = [
-        primitive_op.pipeline if primitive_op is not None else None
-        for primitive_op in predecessor_primitive_ops
-    ]
+    bw_spec = primitive_op.pipeline.config
 
-    # if a predecessor has no primitive op then use 1 for nargs
-    predecessor_funcs_nargs = [
-        pipeline.config.function_nargs if pipeline is not None else 1
-        for pipeline in predecessor_pipelines
+    null_blockwise_spec = BlockwiseSpec(
+        key_function=lambda x: (x,),
+        function=lambda x: x,
+        function_nargs=1,
+        num_input_blocks=(1,),
+        iterable_input_blocks=(False,),
+        reads_map={},
+        writes_list=[],
+    )
+    predecessor_bw_specs = [
+        primitive_op.pipeline.config
+        if primitive_op is not None
+        else null_blockwise_spec
+        for primitive_op in predecessor_primitive_ops
     ]
 
-    mappable = pipeline.mappable
-
-    def apply_pipeline_key_func(pipeline, n_input_blocks, arg):
-        if pipeline is None:
-            return (arg,)
-        if n_input_blocks == 1:
-            assert isinstance(arg, tuple)
-            return pipeline.config.key_function(arg)
-        else:
-            # more than one input block is being read from arg
-            assert isinstance(arg, (list, Iterator))
-            if isinstance(arg, list):
-                return tuple(
-                    list(item)
-                    for item in zip(*(pipeline.config.key_function(a) for a in arg))
-                )
-            else:
-                # Return iterators to avoid materializing all array blocks at
-                # once.
-                return tuple(
-                    iter(list(item))
-                    for item in zip(*(pipeline.config.key_function(a) for a in arg))
-                )
-
-    def fused_key_func(out_key):
-        # this will change when multiple outputs are supported
-        args = pipeline.config.key_function(out_key)
-        # split all args to the fused function into groups, one for each predecessor function
-        func_args = tuple(
-            item
-            for i, (p, a) in enumerate(zip(predecessor_pipelines, args))
-            for item in apply_pipeline_key_func(
-                p, pipeline.config.num_input_blocks[i], a
-            )
-        )
-        return split_into(func_args, predecessor_funcs_nargs)
-
-    def apply_pipeline_func(pipeline, n_input_blocks, *args):
-        if pipeline is None:
-            return args[0]
-        if n_input_blocks == 1:
-            ret = pipeline.config.function(*args)
-        else:
-            # More than one input block is being read from this group of args to primitive op.
-            # Note that it is important that a list is not returned to avoid materializing all
-            # array blocks at once.
-            ret = map(lambda item: pipeline.config.function(*item), zip(*args))
-        return ret
-
-    def fused_func_single(*args):
-        # args are grouped appropriately so they can be called by each predecessor function
-        func_args = [
-            apply_pipeline_func(p, pipeline.config.num_input_blocks[i], *a)
-            for i, (p, a) in enumerate(zip(predecessor_pipelines, args))
-        ]
-        return pipeline.config.function(*func_args)
-
-    # multiple outputs
-    def fused_func_generator(*args):
-        # args are grouped appropriately so they can be called by each predecessor function
-        func_args = [
-            apply_pipeline_func(p, pipeline.config.num_input_blocks[i], *a)
-            for i, (p, a) in enumerate(zip(predecessor_pipelines, args))
-        ]
-        yield from pipeline.config.function(*func_args)
-
-    fused_func = (
-        fused_func_generator
-        if inspect.isgeneratorfunction(pipeline.config.function)
-        else fused_func_single
-    )
-    fused_function_nargs = pipeline.config.function_nargs
-    # ok to get num_input_blocks[0] since it is uniform (see check in can_fuse_multiple_primitive_ops)
-    fused_num_input_blocks = tuple(
-        pipeline.config.num_input_blocks[0] * n
-        for n in itertools.chain(
-            *(
-                p.pipeline.config.num_input_blocks if p is not None else (1,)
-                for p in predecessor_primitive_ops
-            )
-        )
-    )
-    read_proxies = dict(pipeline.config.reads_map)
-    for p in predecessor_pipelines:
-        if p is not None:
-            read_proxies.update(p.config.reads_map)
-    write_proxies = pipeline.config.writes_list
-    spec = BlockwiseSpec(
-        fused_key_func,
-        fused_func,
-        fused_function_nargs,
-        fused_num_input_blocks,
-        read_proxies,
-        write_proxies,
-    )
+    spec = fuse_blockwise_specs(bw_spec, *predecessor_bw_specs)
 
     source_array_names = []
     for i, p in enumerate(predecessor_primitive_ops):
@@ -693,7 +604,7 @@ def fuse_multiple(
     fused_pipeline = CubedPipeline(
         apply_blockwise,
         gensym("fused_apply_blockwise"),
-        mappable,
+        primitive_op.pipeline.mappable,
         spec,
     )
     return PrimitiveOperation(
@@ -708,6 +619,130 @@ def fuse_multiple(
     )
 
 
+def fuse_blockwise_specs(
+    bw_spec: BlockwiseSpec, *predecessor_bw_specs: BlockwiseSpec
+) -> BlockwiseSpec:
+    """
+    Fuse a blockwise spec and its predecessors into a single spec.
+    """
+
+    predecessor_funcs_nargs = [bws.function_nargs for bws in predecessor_bw_specs]
+
+    fused_key_func = make_fused_key_function(
+        bw_spec.key_function,
+        [bws.key_function for bws in predecessor_bw_specs],
+        predecessor_funcs_nargs,
+    )
+
+    fused_func = make_fused_function(
+        bw_spec.function,
+        [bws.function for bws in predecessor_bw_specs],
+        bw_spec.iterable_input_blocks,
+    )
+
+    fused_function_nargs = bw_spec.function_nargs
+    num_input_blocks = bw_spec.num_input_blocks
+    predecessor_num_blocks = [bws.num_input_blocks for bws in predecessor_bw_specs]
+    # note that the length of fused_num_input_blocks is the same as the number of input arrays
+    # but may be different to fused_function_nargs since the fused function groups args
+    fused_num_input_blocks = tuple(
+        itertools.chain(
+            *(
+                tuple(n * m for m in nb)
+                for n, nb in zip(num_input_blocks, predecessor_num_blocks)
+            )
+        )
+    )
+    predecessor_iterable_input_blocks = [
+        bws.iterable_input_blocks for bws in predecessor_bw_specs
+    ]
+    fused_iterable_input_blocks = tuple(
+        itertools.chain(*predecessor_iterable_input_blocks)
+    )
+
+    read_proxies = dict(bw_spec.reads_map)
+    for bws in predecessor_bw_specs:
+        read_proxies.update(bws.reads_map)
+    write_proxies = bw_spec.writes_list
+    return BlockwiseSpec(
+        fused_key_func,
+        fused_func,
+        fused_function_nargs,
+        fused_num_input_blocks,
+        fused_iterable_input_blocks,
+        read_proxies,
+        write_proxies,
+    )
+
+
+def apply_blockwise_key_func(key_function, arg):
+    if isinstance(arg, tuple):
+        return key_function(arg)
+    else:
+        # more than one input block is being read from arg
+        assert isinstance(arg, (list, Iterator))
+        if isinstance(arg, list):
+            return tuple(list(item) for item in zip(*(key_function(a) for a in arg)))
+        else:
+            # Return iterators to avoid materializing all array blocks at
+            # once.
+            return tuple(
+                iter(list(item)) for item in zip(*(key_function(a) for a in arg))
+            )
+
+
+def apply_blockwise_func(func, is_iterable, *args):
+    if is_iterable is False:
+        ret = func(*args)
+    else:
+        # More than one input block is being read from this group of args to primitive op.
+        # Note that it is important that a list is not returned to avoid materializing all
+        # array blocks at once.
+        ret = map(lambda item: func(*item), zip(*args))
+    return ret
+
+
+def make_fused_key_function(
+    key_function, predecessor_key_functions, predecessor_funcs_nargs
+):
+    def fused_key_func(out_key):
+        args = key_function(out_key)
+        # split all args to the fused function into groups, one for each predecessor function
+        func_args = tuple(
+            item
+            for pkf, a in zip(predecessor_key_functions, args)
+            for item in apply_blockwise_key_func(pkf, a)
+        )
+        return split_into(func_args, predecessor_funcs_nargs)
+
+    return fused_key_func
+
+
+def make_fused_function(function, predecessor_functions, iterable_input_blocks):
+    def fused_func_single(*args):
+        # args are grouped appropriately so they can be called by each predecessor function
+        func_args = [
+            apply_blockwise_func(pf, iterable_input_blocks[i], *a)
+            for i, (pf, a) in enumerate(zip(predecessor_functions, args))
+        ]
+        return function(*func_args)
+
+    # multiple outputs
+    def fused_func_generator(*args):
+        # args are grouped appropriately so they can be called by each predecessor function
+        func_args = [
+            apply_blockwise_func(pf, iterable_input_blocks[i], *a)
+            for i, (pf, a) in enumerate(zip(predecessor_functions, args))
+        ]
+        yield from function(*func_args)
+
+    return (
+        fused_func_generator
+        if inspect.isgeneratorfunction(function)
+        else fused_func_single
+    )
+
+
 # blockwise key functions
 
 


Test patch:
diff --git a/cubed/tests/primitive/test_blockwise_fusion.py b/cubed/tests/primitive/test_blockwise_fusion.py
new file mode 100644
index 00000000..f92fc888
--- /dev/null
+++ b/cubed/tests/primitive/test_blockwise_fusion.py
@@ -0,0 +1,302 @@
+import collections.abc
+from itertools import product
+from typing import Iterator
+
+from cubed.primitive.blockwise import (
+    BlockwiseSpec,
+    fuse_blockwise_specs,
+    make_fused_key_function,
+)
+from cubed.utils import map_nested
+
+
+def make_map_blocks_key_function(*names):
+    def key_function(out_key):
+        out_coords = out_key[1:]
+        return tuple((name, *out_coords) for name in names)
+
+    return key_function
+
+
+def make_combine_blocks_list_key_function(name, numblocks, split_every):
+    # similar to the key function in partial_reduce
+    def key_function(out_key):
+        out_coords = out_key[1:]
+
+        # return a tuple with a single item that is a list of input keys to be combined
+        in_keys = [
+            list(
+                range(
+                    bi * split_every,
+                    min((bi + 1) * split_every, numblocks),
+                )
+            )
+            for bi in out_coords
+        ]
+        return ([(name,) + tuple(p) for p in product(*in_keys)],)
+
+    return key_function
+
+
+def make_combine_blocks_iter_key_function(name, numblocks, split_every):
+    # similar to the key function in partial_reduce
+    def key_function(out_key):
+        out_coords = out_key[1:]
+
+        # return a tuple with a single item that is an iterator of input keys to be combined
+        in_keys = [
+            list(
+                range(
+                    bi * split_every,
+                    min((bi + 1) * split_every, numblocks),
+                )
+            )
+            for bi in out_coords
+        ]
+        return (iter([(name,) + tuple(p) for p in product(*in_keys)]),)
+
+    return key_function
+
+
+def negative(x):
+    assert isinstance(x, int)
+    return -x
+
+
+def add(x, y):
+    assert isinstance(x, int)
+    assert isinstance(y, int)
+    return x + y
+
+
+def sum_iter(x):
+    assert isinstance(x, Iterator)
+    return sum(a for a in x)
+
+
+def check_key_function(key_function, out_coords, expected_str):
+    res = key_function(("out",) + out_coords)
+
+    assert str(tuple(iter_repr_nested(r) for r in res)) == expected_str
+
+
+def iter_repr_nested(seq):
+    # convert nested iterators to lists
+    if isinstance(seq, list):
+        return [iter_repr_nested(item) for item in seq]
+    elif isinstance(seq, Iterator):
+        return IteratorWithRepr([iter_repr_nested(item) for item in seq])
+    else:
+        return seq
+
+
+class IteratorWithRepr(collections.abc.Iterator):
+    def __init__(self, values):
+        self.values = values
+        self.it = iter(values)
+
+    def __next__(self):
+        return next(self.it)
+
+    def __repr__(self):
+        return "<" + ", ".join(repr(v) for v in self.values) + ">"
+
+    def __str__(self):
+        return "<" + ", ".join(str(v) for v in self.values) + ">"
+
+
+def make_blockwise_spec(
+    key_function,
+    function,
+    function_nargs=1,
+    num_input_blocks=(1,),
+    iterable_input_blocks=(False,),
+):
+    return BlockwiseSpec(
+        key_function=key_function,
+        function=function,
+        function_nargs=function_nargs,
+        num_input_blocks=num_input_blocks,
+        iterable_input_blocks=iterable_input_blocks,
+        reads_map={},  # unused
+        writes_list=[],  # unused
+    )
+
+
+def test_map_blocks_key_function():
+    key_function = make_map_blocks_key_function("a")
+
+    check_key_function(key_function, (0,), "(('a', 0),)")
+    check_key_function(key_function, (1,), "(('a', 1),)")
+
+
+def test_map_blocks_multiple_inputs_key_function():
+    key_function = make_map_blocks_key_function("a", "b")
+
+    check_key_function(key_function, (0,), "(('a', 0), ('b', 0))")
+    check_key_function(key_function, (1,), "(('a', 1), ('b', 1))")
+
+
+def test_combine_blocks_list_key_function():
+    key_function = make_combine_blocks_list_key_function(
+        "a", numblocks=5, split_every=2
+    )
+
+    check_key_function(key_function, (0,), "([('a', 0), ('a', 1)],)")
+    check_key_function(key_function, (1,), "([('a', 2), ('a', 3)],)")
+    check_key_function(key_function, (2,), "([('a', 4)],)")
+
+
+def test_combine_blocks_iter_key_function():
+    key_function = make_combine_blocks_iter_key_function(
+        "a", numblocks=5, split_every=2
+    )
+
+    check_key_function(key_function, (0,), "(<('a', 0), ('a', 1)>,)")
+    check_key_function(key_function, (1,), "(<('a', 2), ('a', 3)>,)")
+    check_key_function(key_function, (2,), "(<('a', 4)>,)")
+
+
+def test_fuse_key_function_single_multiple():
+    key_function1 = make_map_blocks_key_function("a")
+    key_function2 = make_combine_blocks_iter_key_function(
+        "b", numblocks=5, split_every=2
+    )
+    fused_key_function = make_fused_key_function(key_function2, [key_function1], [1])
+
+    check_key_function(fused_key_function, (0,), "([<('a', 0), ('a', 1)>],)")
+    check_key_function(fused_key_function, (1,), "([<('a', 2), ('a', 3)>],)")
+    check_key_function(fused_key_function, (2,), "([<('a', 4)>],)")
+
+
+def test_fuse_key_function_multiple_single():
+    key_function1 = make_combine_blocks_iter_key_function(
+        "a", numblocks=5, split_every=2
+    )
+    key_function2 = make_map_blocks_key_function("b")
+    fused_key_function = make_fused_key_function(key_function2, [key_function1], [1])
+
+    check_key_function(fused_key_function, (0,), "([<('a', 0), ('a', 1)>],)")
+    check_key_function(fused_key_function, (1,), "([<('a', 2), ('a', 3)>],)")
+    check_key_function(fused_key_function, (2,), "([<('a', 4)>],)")
+
+
+def test_fuse_key_function_multiple_multiple():
+    key_function1 = make_combine_blocks_iter_key_function(
+        "a", numblocks=5, split_every=2
+    )
+    key_function2 = make_combine_blocks_iter_key_function(
+        "b", numblocks=3, split_every=2
+    )
+    fused_key_function = make_fused_key_function(key_function2, [key_function1], [1])
+
+    check_key_function(
+        fused_key_function, (0,), "([<<('a', 0), ('a', 1)>, <('a', 2), ('a', 3)>>],)"
+    )
+    check_key_function(fused_key_function, (1,), "([<<('a', 4)>>],)")
+
+
+def apply_blockwise(input_data, out_coords, bw_spec):
+    args = []
+    out_key = ("out",) + tuple(out_coords)  # array name is ignored by key_function
+    in_keys = bw_spec.key_function(out_key)
+    for in_key in in_keys:
+        # just return the (1D) coord as a value
+        def get_data(key):
+            name = key[0]
+            index = key[1]  # 1d index
+            return input_data[name][index]
+
+        arg = map_nested(get_data, in_key)
+        args.append(arg)
+    return bw_spec.function(*args)
+
+
+def test_apply_blockwise():
+    bw_spec = make_blockwise_spec(
+        key_function=make_map_blocks_key_function("a"),
+        function=negative,
+    )
+
+    input_data = {"a": [0, 1, 2, 3, 4]}
+    out = [apply_blockwise(input_data, [i], bw_spec) for i in range(5)]
+    assert out == [0, -1, -2, -3, -4]
+
+
+def test_apply_blockwise_multiple_inputs():
+    bw_spec = make_blockwise_spec(
+        key_function=make_map_blocks_key_function("a", "b"), function=add
+    )
+
+    input_data = {"a": [0, 1, 2, 3, 4], "b": [5, 6, 7, 8, 9]}
+    out = [apply_blockwise(input_data, [i], bw_spec) for i in range(5)]
+    assert out == [5, 7, 9, 11, 13]
+
+
+def test_apply_blockwise_iterator():
+    bw_spec = make_blockwise_spec(
+        key_function=make_combine_blocks_iter_key_function(
+            "a", numblocks=5, split_every=2
+        ),
+        function=sum_iter,
+    )
+
+    input_data = {"a": [0, 1, 2, 3, 4]}
+    out = [apply_blockwise(input_data, [i], bw_spec) for i in range(3)]
+    assert out == [1, 5, 4]
+
+
+def test_apply_blockwise_fused():
+    bw_spec1 = make_blockwise_spec(
+        key_function=make_map_blocks_key_function("a"), function=negative
+    )
+    bw_spec2 = make_blockwise_spec(
+        key_function=make_map_blocks_key_function("b"), function=negative
+    )
+
+    bw_spec = fuse_blockwise_specs(bw_spec2, bw_spec1)
+
+    input_data = {"a": [0, 1, 2, 3, 4]}
+    out = [apply_blockwise(input_data, [i], bw_spec) for i in range(5)]
+    assert out == [0, 1, 2, 3, 4]
+
+
+def test_apply_blockwise_fused_iterator():
+    bw_spec1 = make_blockwise_spec(
+        key_function=make_map_blocks_key_function("a"), function=negative
+    )
+    bw_spec2 = make_blockwise_spec(
+        key_function=make_combine_blocks_iter_key_function(
+            "b", numblocks=5, split_every=2
+        ),
+        function=sum_iter,
+        num_input_blocks=(2,),
+        iterable_input_blocks=(True,),
+    )
+
+    bw_spec = fuse_blockwise_specs(bw_spec2, bw_spec1)
+
+    input_data = {"a": [0, 1, 2, 3, 4]}
+    out = [apply_blockwise(input_data, [i], bw_spec) for i in range(3)]
+    assert out == [-1, -5, -4]
+
+
+def test_apply_blockwise_fused_iterator_with_single_input_block():
+    bw_spec1 = make_blockwise_spec(
+        key_function=make_map_blocks_key_function("a"), function=negative
+    )
+    # the iterator only reads from a single input block
+    bw_spec2 = make_blockwise_spec(
+        key_function=make_combine_blocks_iter_key_function(
+            "b", numblocks=5, split_every=1
+        ),
+        function=sum_iter,
+        num_input_blocks=(1,),
+        iterable_input_blocks=(True,),
+    )
+
+    bw_spec = fuse_blockwise_specs(bw_spec2, bw_spec1)
+
+    input_data = {"a": [0, 1, 2, 3, 4]}
+    out = [apply_blockwise(input_data, [i], bw_spec) for i in range(5)]
+    assert out == [0, -1, -2, -3, -4]
diff --git a/cubed/tests/test_array_api.py b/cubed/tests/test_array_api.py
index 9ba0a36c..9c09d667 100644
--- a/cubed/tests/test_array_api.py
+++ b/cubed/tests/test_array_api.py
@@ -517,11 +517,13 @@ def test_expand_dims(spec, executor):
     [
         ((10,), (4,), None),
         ((10,), (4,), 0),
+        ((10,), (5,), 0),
         ((10, 7), (4, 3), None),
         ((10, 7), (4, 3), 0),
         ((10, 7), (4, 3), 1),
         ((10, 7), (4, 3), (0, 1)),
         ((10, 7), (4, 3), -1),
+        ((10, 7), (5, 3), (0, 1)),
     ],
 )
 def test_flip(executor, shape, chunks, axis):
diff --git a/cubed/tests/test_groupby.py b/cubed/tests/test_groupby.py
index 10bff068..ab047893 100644
--- a/cubed/tests/test_groupby.py
+++ b/cubed/tests/test_groupby.py
@@ -96,7 +96,6 @@ def test_get_chunks_for_groups(
 def test_groupby_blockwise_axis0():
     a = xp.ones((10, 3), dtype=nxp.int32, chunks=(6, 2))
     b = nxp.asarray([0, 0, 0, 1, 1, 2, 2, 4, 4, 4])
-    extra_func_kwargs = dict(dtype=nxp.int32)
     c = groupby_blockwise(
         a,
         b,
@@ -104,7 +103,6 @@ def test_groupby_blockwise_axis0():
         axis=0,
         dtype=nxp.int64,
         num_groups=6,
-        extra_func_kwargs=extra_func_kwargs,
     )
     assert_array_equal(
         c.compute(),
@@ -124,7 +122,6 @@ def test_groupby_blockwise_axis0():
 def test_groupby_blockwise_axis1():
     a = xp.ones((3, 10), dtype=nxp.int32, chunks=(6, 2))
     b = nxp.asarray([0, 0, 0, 1, 1, 2, 2, 4, 4, 4])
-    extra_func_kwargs = dict(dtype=nxp.int32)
     c = groupby_blockwise(
         a,
         b,
@@ -132,7 +129,6 @@ def test_groupby_blockwise_axis1():
         axis=1,
         dtype=nxp.int64,
         num_groups=6,
-        extra_func_kwargs=extra_func_kwargs,
     )
     assert_array_equal(
         c.compute(),
@@ -146,7 +142,7 @@ def test_groupby_blockwise_axis1():
     )
 
 
-def _sum_reduction_func(arr, by, axis, start_group, num_groups, dtype):
+def _sum_reduction_func(arr, by, axis, start_group, num_groups):
     # change 'by' so it starts from 0 for each chunk
     by = by - start_group
-    return npg.aggregate(by, arr, func="sum", dtype=dtype, axis=axis, size=num_groups)
+    return npg.aggregate(by, arr, func="sum", axis=axis, size=num_groups)
diff --git a/cubed/tests/test_mem_utilization.py b/cubed/tests/test_mem_utilization.py
index 8332b45a..a384b90f 100644
--- a/cubed/tests/test_mem_utilization.py
+++ b/cubed/tests/test_mem_utilization.py
@@ -93,6 +93,15 @@ def test_index_chunk_aligned(tmp_path, spec, executor):
     run_operation(tmp_path, executor, "index_chunk_aligned", b)
 
 
+@pytest.mark.slow
+def test_index_multiple_axes(tmp_path, spec, executor):
+    a = cubed.random.random(
+        (10000, 10000), chunks=(5000, 5000), spec=spec
+    )  # 200MB chunks
+    b = a[1:, 1:]
+    run_operation(tmp_path, executor, "index_multiple_axes", b)
+
+
 @pytest.mark.slow
 def test_index_step(tmp_path, spec, executor):
     a = cubed.random.random(
@@ -265,6 +274,16 @@ def test_flip(tmp_path, spec, executor):
     run_operation(tmp_path, executor, "flip", b)
 
 
+@pytest.mark.slow
+def test_flip_multiple_axes(tmp_path, spec, executor):
+    # Note 'a' has one fewer element in both axes to force chunking to cross array boundaries
+    a = cubed.random.random(
+        (9999, 9999), chunks=(5000, 5000), spec=spec
+    )  # 200MB chunks
+    b = xp.flip(a)
+    run_operation(tmp_path, executor, "flip_multiple_axes", b)
+
+
 @pytest.mark.slow
 def test_reshape(tmp_path, spec, executor):
     a = cubed.random.random(