Repository: cubed-dev/cubed (python)

Issue:
Multi-stage rechunk
The rechunker algorithm has support for [multi-stage rechunking](https://github.com/pangeo-data/rechunker/pull/89), which reduces the number of IO operations.

It would be good to enable this in Cubed `rechunk`. The simplest way to do this would be to expose the `min_mem` parameter, but it might be interesting to see if there is a way of automatically finding the number of stages that optimizes for something we care about (e.g. min IO ops).

Pull request description:
Multistage rechunking
Fixes #650

Patch:
diff --git a/cubed/array_api/manipulation_functions.py b/cubed/array_api/manipulation_functions.py
index aa3682cb..9a669300 100644
--- a/cubed/array_api/manipulation_functions.py
+++ b/cubed/array_api/manipulation_functions.py
@@ -167,7 +167,7 @@ def concat(arrays, /, *, axis=0, chunks=None):
     num_input_blocks = (1,) * len(arrays)
     iterable_input_blocks = (True,) * len(arrays)
 
-    # We have to mark this as fusable=False since the number of input args to
+    # We have to mark this as fusable_with_predecessors=False since the number of input args to
     # the _read_concat_chunk function is *not* the same as the number of
     # predecessor nodes in the DAG, and the fusion functions in blockwise
     # assume they are the same. See https://github.com/cubed-dev/cubed/issues/414
@@ -188,7 +188,7 @@ def concat(arrays, /, *, axis=0, chunks=None):
         offsets=offsets,
         in_shapes=in_shapes,
         function_nargs=1,
-        fusable=False,
+        fusable_with_predecessors=False,
     )
 
 
@@ -559,7 +559,7 @@ def stack(arrays, /, *, axis=0):
         in_name = array_names[out_coords[axis]]
         return ((in_name, *(out_coords[:axis] + out_coords[(axis + 1) :])),)
 
-    # We have to mark this as fusable=False since the number of input args to
+    # We have to mark this as fusable_with_predecessors=False since the number of input args to
     # the _read_stack_chunk function is *not* the same as the number of
     # predecessor nodes in the DAG, and the fusion functions in blockwise
     # assume they are the same. See https://github.com/cubed-dev/cubed/issues/414
@@ -572,7 +572,7 @@ def stack(arrays, /, *, axis=0):
         chunkss=[chunks],
         axis=axis,
         function_nargs=1,
-        fusable=False,
+        fusable_with_predecessors=False,
     )
 
 
diff --git a/cubed/core/array.py b/cubed/core/array.py
index a161927b..b5a61a0e 100644
--- a/cubed/core/array.py
+++ b/cubed/core/array.py
@@ -162,7 +162,9 @@ class CoreArray:
         if result:
             return result[0]
 
-    def rechunk(self: T_ChunkedArray, chunks) -> T_ChunkedArray:
+    def rechunk(
+        self: T_ChunkedArray, chunks, *, min_mem=None, use_new_impl=False
+    ) -> T_ChunkedArray:
         """Change the chunking of this array without changing its shape or data.
 
         Parameters
@@ -177,7 +179,7 @@ class CoreArray:
         """
         from cubed.core.ops import rechunk
 
-        return rechunk(self, chunks)
+        return rechunk(self, chunks, min_mem=min_mem, use_new_impl=use_new_impl)
 
     def visualize(
         self,
diff --git a/cubed/core/ops.py b/cubed/core/ops.py
index 0a3ab7a5..b64a45d6 100644
--- a/cubed/core/ops.py
+++ b/cubed/core/ops.py
@@ -34,6 +34,7 @@ from cubed.vendor.dask.array.core import normalize_chunks
 from cubed.vendor.dask.array.utils import validate_axis
 from cubed.vendor.dask.blockwise import broadcast_dimensions
 from cubed.vendor.dask.utils import has_keyword
+from cubed.vendor.rechunker.algorithm import multistage_rechunking_plan
 
 if TYPE_CHECKING:
     from cubed.array_api.array_object import Array
@@ -285,7 +286,8 @@ def blockwise(
 
     extra_projected_mem = kwargs.pop("extra_projected_mem", 0)
 
-    fusable = kwargs.pop("fusable", True)
+    fusable_with_predecessors = kwargs.pop("fusable_with_predecessors", True)
+    fusable_with_successors = kwargs.pop("fusable_with_successors", True)
     num_input_blocks = kwargs.pop("num_input_blocks", None)
     iterable_input_blocks = kwargs.pop("iterable_input_blocks", None)
 
@@ -311,7 +313,8 @@ def blockwise(
         in_names=in_names,
         out_name=name,
         extra_func_kwargs=extra_func_kwargs,
-        fusable=fusable,
+        fusable_with_predecessors=fusable_with_predecessors,
+        fusable_with_successors=fusable_with_successors,
         num_input_blocks=num_input_blocks,
         iterable_input_blocks=iterable_input_blocks,
         **kwargs,
@@ -436,6 +439,8 @@ def _general_blockwise(
     num_input_blocks = kwargs.pop("num_input_blocks", None)
     iterable_input_blocks = kwargs.pop("iterable_input_blocks", None)
 
+    op_name = kwargs.pop("op_name", "blockwise")
+
     spec = check_array_specs(arrays)
 
     if isinstance(target_stores, list):  # multiple outputs
@@ -471,7 +476,7 @@ def _general_blockwise(
     )
     plan = Plan._new(
         name,
-        "blockwise",
+        op_name,
         op.target_array,
         op,
         False,
@@ -1011,12 +1016,12 @@ def map_direct(
         chunks=chunks,
         extra_source_arrays=args,
         extra_projected_mem=extra_projected_mem,
-        fusable=False,  # don't allow fusion with predecessors since side inputs are not accounted for
+        fusable_with_predecessors=False,  # don't allow fusion with predecessors since side inputs are not accounted for
         **kwargs,
     )
 
 
-def rechunk(x, chunks, target_store=None):
+def rechunk(x, chunks, *, target_store=None, min_mem=None, use_new_impl=False):
     """Change the chunking of an array without changing its shape or data.
 
     Parameters
@@ -1029,6 +1034,9 @@ def rechunk(x, chunks, target_store=None):
     cubed.Array
         An array with the desired chunks.
     """
+    if use_new_impl:
+        return rechunk_new(x, chunks, min_mem=min_mem)
+
     if isinstance(chunks, dict):
         chunks = {validate_axis(c, x.ndim): v for c, v in chunks.items()}
         for i in range(x.ndim):
@@ -1105,6 +1113,166 @@ def rechunk(x, chunks, target_store=None):
         return Array(name, op2.target_array, spec, plan2)
 
 
+def rechunk_new(x, chunks, *, min_mem=None):
+    """Change the chunking of an array without changing its shape or data.
+
+    Parameters
+    ----------
+    chunks : tuple
+        The desired chunks of the array after rechunking.
+
+    Returns
+    -------
+    cubed.Array
+        An array with the desired chunks.
+    """
+    if isinstance(chunks, dict):
+        chunks = {validate_axis(c, x.ndim): v for c, v in chunks.items()}
+        for i in range(x.ndim):
+            if i not in chunks:
+                chunks[i] = x.chunks[i]
+            elif chunks[i] is None:
+                chunks[i] = x.chunks[i]
+    if isinstance(chunks, (tuple, list)):
+        chunks = tuple(lc if lc is not None else rc for lc, rc in zip(chunks, x.chunks))
+
+    normalized_chunks = normalize_chunks(chunks, x.shape, dtype=x.dtype)
+    if x.chunks == normalized_chunks:
+        return x
+    # normalizing takes care of dict args for chunks
+    target_chunks = to_chunksize(normalized_chunks)
+
+    # merge chunks special case
+    if all(c1 % c0 == 0 for c0, c1 in zip(x.chunksize, target_chunks)):
+        return merge_chunks(x, target_chunks)
+
+    spec = x.spec
+    source_chunks = to_chunksize(normalize_chunks(x.chunks, x.shape, dtype=x.dtype))
+
+    # rechunker doesn't take account of uncompressed and compressed copies of the
+    # input and output array chunk/selection, so adjust appropriately
+    # note the factor is 5 (not 4) since there is a extra (unnecessary) copy
+    # made when writing out to Zarr
+    rechunker_max_mem = (spec.allowed_mem - spec.reserved_mem) // 5
+    if min_mem is None:
+        min_mem = min(rechunker_max_mem // 20, x.nbytes)
+    stages = multistage_rechunking_plan(
+        shape=x.shape,
+        source_chunks=source_chunks,
+        target_chunks=target_chunks,
+        itemsize=x.dtype.itemsize,
+        min_mem=min_mem,
+        max_mem=rechunker_max_mem,
+    )
+
+    out = x
+    for i, stage in enumerate(stages):
+        last_stage = i == len(stages) - 1
+        read_chunks, int_chunks, write_chunks = stage
+
+        # Use target chunks for last stage
+        target_chunks_ = target_chunks if last_stage else write_chunks
+
+        if read_chunks == write_chunks:
+            out = _rechunk(out, read_chunks, target_chunks_)
+        else:
+            intermediate = _rechunk(out, read_chunks, int_chunks)
+            out = _rechunk(intermediate, write_chunks, target_chunks_)
+
+    return out
+
+
+def _rechunk(x, copy_chunks, target_chunks):
+    # rechunk x so that its target store has target_chunks, using copy_chunks as the size of chunks for copying from source to target
+
+    normalized_copy_chunks = normalize_chunks(copy_chunks, x.shape, dtype=x.dtype)
+    copy_chunks = to_chunksize(normalized_copy_chunks)
+
+    copy_chunks_mem = array_memory(x.dtype, copy_chunks)
+
+    target_chunks = normalize_chunks(target_chunks, x.shape, dtype=x.dtype)
+    target_chunks = to_chunksize(target_chunks)
+
+    def selection_function(out_key):
+        out_coords = out_key[1:]
+        return get_item(normalized_copy_chunks, out_coords)
+
+    max_num_input_blocks = math.prod(
+        math.ceil(c1 / c0) for c0, c1 in zip(x.chunksize, copy_chunks)
+    )
+
+    return map_selection(
+        None,  # no function to apply after selection
+        selection_function,
+        x,
+        x.shape,
+        x.dtype,
+        normalized_copy_chunks,
+        max_num_input_blocks=max_num_input_blocks,
+        target_chunks_=target_chunks,
+        fusable_with_predecessors=False,
+        fusable_with_successors=False,
+        op_name="rechunk",
+        extra_projected_mem=copy_chunks_mem,
+    )
+
+
+def rechunk_plan(x, chunks, *, min_mem=None):
+    if isinstance(chunks, dict):
+        chunks = {validate_axis(c, x.ndim): v for c, v in chunks.items()}
+        for i in range(x.ndim):
+            if i not in chunks:
+                chunks[i] = x.chunks[i]
+            elif chunks[i] is None:
+                chunks[i] = x.chunks[i]
+    if isinstance(chunks, (tuple, list)):
+        chunks = tuple(lc if lc is not None else rc for lc, rc in zip(chunks, x.chunks))
+
+    normalized_chunks = normalize_chunks(chunks, x.shape, dtype=x.dtype)
+    if x.chunks == normalized_chunks:
+        return x
+    # normalizing takes care of dict args for chunks
+    target_chunks = to_chunksize(normalized_chunks)
+
+    # merge chunks special case
+    if all(c1 % c0 == 0 for c0, c1 in zip(x.chunksize, target_chunks)):
+        return merge_chunks(x, target_chunks)
+
+    spec = x.spec
+    source_chunks = to_chunksize(normalize_chunks(x.chunks, x.shape, dtype=x.dtype))
+
+    # rechunker doesn't take account of uncompressed and compressed copies of the
+    # input and output array chunk/selection, so adjust appropriately
+    rechunker_max_mem = (spec.allowed_mem - spec.reserved_mem) // 5
+    if min_mem is None:
+        min_mem = min(rechunker_max_mem // 20, x.nbytes)
+    stages = multistage_rechunking_plan(
+        shape=x.shape,
+        source_chunks=source_chunks,
+        target_chunks=target_chunks,
+        itemsize=x.dtype.itemsize,
+        min_mem=min_mem,
+        max_mem=rechunker_max_mem,
+    )
+
+    source_chunks = x.chunksize
+    for i, stage in enumerate(stages):
+        last_stage = i == len(stages) - 1
+        read_chunks, int_chunks, write_chunks = stage
+
+        # Use target chunks for last stage
+        target_chunks_ = target_chunks if last_stage else write_chunks
+
+        if read_chunks == write_chunks:
+            yield source_chunks, read_chunks, target_chunks_
+            source_chunks = target_chunks_
+        else:
+            yield source_chunks, read_chunks, int_chunks
+            source_chunks = int_chunks
+            yield source_chunks, write_chunks, target_chunks_
+            source_chunks = target_chunks_
+
+
 def merge_chunks(x, chunks):
     """Merge multiple chunks into one."""
     target_chunksize = chunks
diff --git a/cubed/core/optimization.py b/cubed/core/optimization.py
index 5ec5064c..1c5af647 100644
--- a/cubed/core/optimization.py
+++ b/cubed/core/optimization.py
@@ -134,9 +134,12 @@ def is_primitive_op(node_dict):
     return "primitive_op" in node_dict
 
 
-def is_fusable(node_dict):
+def is_fusable_with_predecessors(node_dict):
     """Return True if a node is a primitive op and can be fused with its predecessors."""
-    return is_primitive_op(node_dict) and node_dict["primitive_op"].fusable
+    return (
+        is_primitive_op(node_dict)
+        and node_dict["primitive_op"].fusable_with_predecessors
+    )
 
 
 def num_source_arrays(dag, name):
@@ -165,7 +168,7 @@ def can_fuse_predecessors(
     nodes = dict(dag.nodes(data=True))
 
     # if node itself can't be fused then there is nothing to fuse
-    if not is_fusable(nodes[name]):
+    if not is_fusable_with_predecessors(nodes[name]):
         logger.debug(
             "can't fuse %s since it is not a primitive operation, or it uses an operation that can't be fused (concat or stack)",
             name,
diff --git a/cubed/core/plan.py b/cubed/core/plan.py
index 8e54993b..7ea6499c 100644
--- a/cubed/core/plan.py
+++ b/cubed/core/plan.py
@@ -408,6 +408,9 @@ class Plan:
                         tooltip += (
                             f"\nnum input blocks: {pipeline.config.num_input_blocks}"
                         )
+                        tooltip += (
+                            f"\nnum output blocks: {pipeline.config.num_output_blocks}"
+                        )
                     del d["pipeline"]
 
                 if "stack_summaries" in d and d["stack_summaries"] is not None:
@@ -584,5 +587,5 @@ def create_zarr_arrays(lazy_zarr_arrays, allowed_mem, reserved_mem):
         allowed_mem=allowed_mem,
         reserved_mem=reserved_mem,
         num_tasks=num_tasks,
-        fusable=False,
+        fusable_with_predecessors=False,
     )
diff --git a/cubed/primitive/blockwise.py b/cubed/primitive/blockwise.py
index f133d7f1..628d1fd1 100644
--- a/cubed/primitive/blockwise.py
+++ b/cubed/primitive/blockwise.py
@@ -19,7 +19,7 @@ from cubed.primitive.memory import BufferCopies, MemoryModeller, calculate_proje
 from cubed.primitive.types import CubedArrayProxy, PrimitiveOperation
 from cubed.runtime.types import CubedPipeline
 from cubed.storage.zarr import LazyZarrArray, T_ZarrArray, lazy_zarr_array
-from cubed.types import T_Chunks, T_DType, T_Shape, T_Store
+from cubed.types import T_Chunks, T_DType, T_RegularChunks, T_Shape, T_Store
 from cubed.utils import array_memory, chunk_memory, get_item, map_nested
 from cubed.utils import numblocks as compute_numblocks
 from cubed.utils import split_into, to_chunksize
@@ -57,6 +57,8 @@ class BlockwiseSpec:
         arrays that the operation depends on.
     num_input_blocks: Tuple[int, ...]
         The number of input blocks read from each input array.
+    num_output_blocks: Tuple[int, ...]
+        The number of output blocks written to each output 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]
@@ -69,6 +71,7 @@ class BlockwiseSpec:
     function: Callable[..., Any]
     function_nargs: int
     num_input_blocks: Tuple[int, ...]
+    num_output_blocks: Tuple[int, ...]
     iterable_input_blocks: Tuple[bool, ...]
     reads_map: Dict[str, CubedArrayProxy]
     writes_list: List[CubedArrayProxy]
@@ -160,7 +163,8 @@ def blockwise(
     out_name: Optional[str] = None,
     extra_projected_mem: int = 0,
     extra_func_kwargs: Optional[Dict[str, Any]] = None,
-    fusable: bool = True,
+    fusable_with_predecessors: bool = True,
+    fusable_with_successors: bool = True,
     num_input_blocks: Optional[Tuple[int, ...]] = None,
     iterable_input_blocks: Optional[Tuple[bool, ...]] = None,
     **kwargs,
@@ -247,7 +251,8 @@ def blockwise(
         in_names=in_names,
         extra_projected_mem=extra_projected_mem,
         extra_func_kwargs=extra_func_kwargs,
-        fusable=fusable,
+        fusable_with_predecessors=fusable_with_predecessors,
+        fusable_with_successors=fusable_with_successors,
         num_input_blocks=num_input_blocks,
         iterable_input_blocks=iterable_input_blocks,
         **kwargs,
@@ -270,10 +275,12 @@ def general_blockwise(
     in_names: Optional[List[str]] = None,
     extra_projected_mem: int = 0,
     extra_func_kwargs: Optional[Dict[str, Any]] = None,
-    fusable: bool = True,
+    fusable_with_predecessors: bool = True,
+    fusable_with_successors: bool = True,
     function_nargs: Optional[int] = None,
     num_input_blocks: Optional[Tuple[int, ...]] = None,
     iterable_input_blocks: Optional[Tuple[bool, ...]] = None,
+    target_chunks_: Optional[T_RegularChunks] = None,
     return_writes_stores: bool = False,
     **kwargs,
 ) -> PrimitiveOperation:
@@ -351,7 +358,7 @@ def general_blockwise(
                 target_store,
                 shapes[i],
                 dtype=dtypes[i],
-                chunks=chunksize,
+                chunks=target_chunks_ or chunksize,
                 path=target_paths[i] if target_paths is not None else None,
                 storage_options=storage_options,
                 compressor=compressor,
@@ -365,11 +372,16 @@ def general_blockwise(
             output_chunk_memory, array_memory(dtypes[i], chunksize)
         )
 
+    # the number of blocks written to each target array is currently the same
+    nb = math.prod(chunksize) // math.prod(target_chunks_ or chunksize)
+    num_output_blocks = (nb,) * len(target_arrays)
+
     spec = BlockwiseSpec(
         key_function,
         func_with_kwargs,
         function_nargs,
         num_input_blocks,
+        num_output_blocks,
         iterable_input_blocks,
         read_proxies,
         write_proxies,
@@ -411,7 +423,9 @@ def general_blockwise(
         allowed_mem=allowed_mem,
         reserved_mem=reserved_mem,
         num_tasks=num_tasks,
-        fusable=fusable,
+        fusable_with_predecessors=fusable_with_predecessors,
+        fusable_with_successors=fusable_with_successors,
+        write_chunks=chunksize,
     )
 
 
@@ -422,7 +436,9 @@ def is_fuse_candidate(primitive_op: PrimitiveOperation) -> bool:
     """
     Return True if a primitive operation is a candidate for blockwise fusion.
     """
-    return primitive_op.pipeline.function == apply_blockwise
+    return primitive_op.pipeline.function == apply_blockwise and (
+        primitive_op.fusable_with_predecessors or primitive_op.fusable_with_successors
+    )
 
 
 def can_fuse_primitive_ops(
@@ -543,12 +559,14 @@ def fuse(
         n * pipeline2.config.num_input_blocks[0]
         for n in pipeline1.config.num_input_blocks
     )
+    num_output_blocks = pipeline2.config.num_output_blocks
     iterable_input_blocks = pipeline1.config.iterable_input_blocks
     spec = BlockwiseSpec(
         fused_key_func,
         fused_func,
         function_nargs,
         num_input_blocks,
+        num_output_blocks,
         iterable_input_blocks,
         read_proxies,
         write_proxies,
@@ -576,7 +594,7 @@ def fuse(
         allowed_mem=allowed_mem,
         reserved_mem=reserved_mem,
         num_tasks=num_tasks,
-        fusable=True,
+        fusable_with_predecessors=True,
     )
 
 
@@ -594,6 +612,7 @@ def fuse_multiple(
         function=lambda x: x,
         function_nargs=1,
         num_input_blocks=(1,),
+        num_output_blocks=(1,),
         iterable_input_blocks=(False,),
         reads_map={},
         writes_list=[],
@@ -636,7 +655,7 @@ def fuse_multiple(
         allowed_mem=allowed_mem,
         reserved_mem=reserved_mem,
         num_tasks=num_tasks,
-        fusable=True,
+        fusable_with_predecessors=True,
     )
 
 
@@ -674,6 +693,7 @@ def fuse_blockwise_specs(
             )
         )
     )
+    fused_num_output_blocks = bw_spec.num_output_blocks
     predecessor_iterable_input_blocks = [
         bws.iterable_input_blocks for bws in predecessor_bw_specs
     ]
@@ -691,6 +711,7 @@ def fuse_blockwise_specs(
         fused_func,
         fused_function_nargs,
         fused_num_input_blocks,
+        fused_num_output_blocks,
         fused_iterable_input_blocks,
         read_proxies,
         write_proxies,
diff --git a/cubed/primitive/rechunk.py b/cubed/primitive/rechunk.py
index 06245898..1187cd98 100644
--- a/cubed/primitive/rechunk.py
+++ b/cubed/primitive/rechunk.py
@@ -236,6 +236,6 @@ def spec_to_primitive_op(
         allowed_mem=allowed_mem,
         reserved_mem=reserved_mem,
         num_tasks=num_tasks,
-        fusable=False,
+        fusable_with_predecessors=False,
         write_chunks=spec.write.chunks,
     )
diff --git a/cubed/primitive/types.py b/cubed/primitive/types.py
index 81759fce..6599ace7 100644
--- a/cubed/primitive/types.py
+++ b/cubed/primitive/types.py
@@ -37,9 +37,12 @@ class PrimitiveOperation:
     num_tasks: int
     """The number of tasks needed to run this operation."""
 
-    fusable: bool = True
+    fusable_with_predecessors: bool = True
     """Whether this operation can be fused with predecessor operations."""
 
+    fusable_with_successors: bool = True
+    """Whether this operation can be fused with successor operations."""
+
     write_chunks: Optional[T_RegularChunks] = None
     """The chunk size used by this operation."""
 


Test patch:
diff --git a/cubed/tests/primitive/test_blockwise_fusion.py b/cubed/tests/primitive/test_blockwise_fusion.py
index f92fc888..6aa45f01 100644
--- a/cubed/tests/primitive/test_blockwise_fusion.py
+++ b/cubed/tests/primitive/test_blockwise_fusion.py
@@ -110,6 +110,7 @@ def make_blockwise_spec(
     function,
     function_nargs=1,
     num_input_blocks=(1,),
+    num_output_blocks=(1,),
     iterable_input_blocks=(False,),
 ):
     return BlockwiseSpec(
@@ -117,6 +118,7 @@ def make_blockwise_spec(
         function=function,
         function_nargs=function_nargs,
         num_input_blocks=num_input_blocks,
+        num_output_blocks=num_output_blocks,
         iterable_input_blocks=iterable_input_blocks,
         reads_map={},  # unused
         writes_list=[],  # unused
diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py
index ed17df87..0f9773bc 100644
--- a/cubed/tests/test_core.py
+++ b/cubed/tests/test_core.py
@@ -294,18 +294,19 @@ def test_rechunk_same_chunks(spec):
 
 # see also test_rechunk.py
 def test_rechunk_intermediate(tmp_path):
-    spec = cubed.Spec(tmp_path, allowed_mem=4 * 8 * 4)
-    a = xp.ones((4, 4), chunks=(1, 4), spec=spec)
-    b = a.rechunk((4, 1))
-    assert_array_equal(b.compute(), np.ones((4, 4)))
-    intermediates = [n for (n, d) in b.plan.dag.nodes(data=True) if "-int" in d["name"]]
-    assert len(intermediates) == 1
+    # factor of 4 is for chunks copies, extra 8 is for map_selection
+    spec = cubed.Spec(tmp_path, allowed_mem=5 * 8 * 4 + 8)
+    a = xp.ones((5, 5), chunks=(1, 5), spec=spec)
+    b = a.rechunk((5, 1))
+    assert_array_equal(b.compute(), np.ones((5, 5)))
+    # intermediates = [n for (n, d) in b.plan.dag.nodes(data=True) if "-int" in d["name"]]
+    # assert len(intermediates) == 1
     rechunks = [
         n
         for (n, d) in b.plan.dag.nodes(data=True)
         if d.get("op_name", None) == "rechunk"
     ]
-    assert len(rechunks) > 0
+    assert len(rechunks) == 2  # two ops due to intermediate store
 
 
 def test_rechunk_merge_chunks_optimization():
diff --git a/cubed/tests/test_mem_utilization_rechunk.py b/cubed/tests/test_mem_utilization_rechunk.py
new file mode 100644
index 00000000..4a562f2d
--- /dev/null
+++ b/cubed/tests/test_mem_utilization_rechunk.py
@@ -0,0 +1,88 @@
+import math
+import platform
+import sys
+
+import pytest
+
+import cubed
+import cubed.array_api as xp
+import cubed.random
+from cubed.runtime.create import create_executor
+from cubed.tests.test_mem_utilization import run_operation
+from cubed.tests.utils import LITHOPS_LOCAL_CONFIG
+
+ALLOWED_MEM = 2_500_000_000
+
+EXECUTORS = {}
+
+if platform.system() != "Windows":
+    # Run with max_tasks_per_child=1 so that each task is run in a new process,
+    # allowing us to perform a stronger check on peak memory
+    if sys.version_info >= (3, 11):
+        executor_options = dict(max_tasks_per_child=1, max_workers=4)
+        EXECUTORS["processes-single-task"] = create_executor(
+            "processes", executor_options
+        )
+
+try:
+    executor_options = dict(config=LITHOPS_LOCAL_CONFIG, wait_dur_sec=0.1)
+    EXECUTORS["lithops"] = create_executor("lithops", executor_options)
+except ImportError:
+    pass
+
+
+@pytest.fixture()
+def spec(tmp_path, reserved_mem):
+    return cubed.Spec(tmp_path, allowed_mem=ALLOWED_MEM, reserved_mem=reserved_mem)
+
+
+@pytest.fixture(
+    scope="module",
+    params=EXECUTORS.values(),
+    ids=EXECUTORS.keys(),
+)
+def executor(request):
+    return request.param
+
+
+@pytest.fixture(scope="module")
+def reserved_mem(executor):
+    res = cubed.measure_reserved_mem(executor) * 1.1  # add some wiggle room
+    return round_up_to_multiple(res, 10_000_000)  # round up to nearest multiple of 10MB
+
+
+def round_up_to_multiple(x, multiple=10):
+    """Round up to the nearest multiple"""
+    return math.ceil(x / multiple) * multiple
+
+
+@pytest.mark.slow
+def test_rechunk_era5(tmp_path, spec, executor):
+    # This example is based on rechunking an ERA5 dataset
+    # from https://github.com/pangeo-data/rechunker/pull/89
+    shape = (350640, 721, 1440)
+    source_chunks = (31, 721, 1440)
+    target_chunks = (350640, 10, 10)
+
+    x = cubed.random.random(shape, dtype=xp.float32, chunks=source_chunks, spec=spec)
+
+    from cubed.core.ops import rechunk_plan
+
+    i = 0
+    for source_chunks, copy_chunks, target_chunks in rechunk_plan(x, target_chunks):
+        # Find the smallest shape that contains the three chunk sizes
+        # This will be a lot less than the full ERA5 shape (350640, 721, 1440),
+        # making it suitable for running in a test
+        test_shape = tuple(
+            max(a, b, c) for a, b, c in zip(source_chunks, copy_chunks, target_chunks)
+        )
+        print(i, test_shape, source_chunks, copy_chunks, target_chunks)
+
+        a = cubed.random.random(
+            test_shape, dtype=xp.float32, chunks=source_chunks, spec=spec
+        )
+        b = a.rechunk(target_chunks, use_new_impl=True)
+
+        run_operation(tmp_path, executor, f"rechunk_era5_stage_{i}", b)
+
+        i += 1
diff --git a/cubed/tests/test_optimization.py b/cubed/tests/test_optimization.py
index ebb94022..82f36545 100644
--- a/cubed/tests/test_optimization.py
+++ b/cubed/tests/test_optimization.py
@@ -245,9 +245,9 @@ def create_dag():
     return nx.MultiDiGraph()
 
 
-def add_op(dag, func, inputs, outputs, fusable=True):
+def add_op(dag, func, inputs, outputs, fusable_with_predecessors=True):
     name = gensym(func.__name__)
-    dag.add_node(name, func=func, fusable=fusable)
+    dag.add_node(name, func=func, fusable_with_predecessors=fusable_with_predecessors)
     for n in inputs:
         dag.add_edge(n, name)
     for n in outputs:
@@ -1171,7 +1171,7 @@ def test_fuse_only_optimize_dag(spec):
     d = xp.negative(c)
 
     # only fuse d (with c)
-    # b should remain un-fused, even though it is fusable
+    # b should remain un-fused, even though it is fusable with predecessors
     op_name = next(d.plan.dag.predecessors(d.name))
     opt_fn = partial(fuse_only_optimize_dag, only_fuse=[op_name])
 
@@ -1192,7 +1192,7 @@ def test_fuse_only_optimize_dag(spec):
 
 
 def test_optimize_stack(spec):
-    # This test fails if stack's general_blockwise call doesn't have fusable=False
+    # This test fails if stack's general_blockwise call doesn't have fusable_with_predecessors=False
     a = cubed.random.random((10, 10), chunks=(5, 5), spec=spec)
     b = cubed.random.random((10, 10), chunks=(5, 5), spec=spec)
     c = xp.stack((a, b), axis=0)
@@ -1202,7 +1202,7 @@ def test_optimize_stack(spec):
 
 
 def test_optimize_concat(spec):
-    # This test fails if concat's general_blockwise call doesn't have fusable=False
+    # This test fails if concat's general_blockwise call doesn't have fusable_with_predecessors=False
     a = cubed.random.random((10, 10), chunks=(5, 5), spec=spec)
     b = cubed.random.random((10, 10), chunks=(5, 5), spec=spec)
     c = xp.concat((a, b), axis=0)
diff --git a/cubed/tests/test_rechunk.py b/cubed/tests/test_rechunk.py
new file mode 100644
index 00000000..92fb10dc
--- /dev/null
+++ b/cubed/tests/test_rechunk.py
@@ -0,0 +1,54 @@
+import pytest
+
+import cubed
+import cubed as xp
+
+
+@pytest.mark.parametrize(
+    "min_mem, expected_num_stages, expected_max_input_blocks, expected_max_output_blocks",
+    [
+        [None, 3, 16, 15],  # multistage rechunk - more stages, lower fan in/out
+        [0, 1, 3771, 3460],  # single stage rechunk - very high fan in/out
+    ],
+)
+def test_rechunk_era5(
+    tmp_path,
+    min_mem,
+    expected_num_stages,
+    expected_max_input_blocks,
+    expected_max_output_blocks,
+):
+    # from https://github.com/pangeo-data/rechunker/pull/89
+    shape = (350640, 721, 1440)
+    source_chunks = (31, 721, 1440)
+    target_chunks = (350640, 10, 10)
+
+    spec = cubed.Spec(allowed_mem="2.5GB")
+
+    a = xp.empty(shape, dtype=xp.float32, chunks=source_chunks, spec=spec)
+    b = a.rechunk(target_chunks, min_mem=min_mem, use_new_impl=True)
+
+    b.visualize(filename=tmp_path / "rechunk_era5")
+
+    assert b.shape == a.shape
+
+    # find rechunk ops in plan
+    plan = b.plan._finalize()
+    rechunks = [
+        (n, d)
+        for (n, d) in plan.dag.nodes(data=True)
+        if d.get("op_name", None) == "rechunk"
+    ]
+
+    # each stage has two ops due to intermediate store
+    assert len(rechunks) == expected_num_stages * 2
+
+    max_input_blocks = max(
+        d["pipeline"].config.num_input_blocks[0] for _, d in rechunks
+    )
+    assert max_input_blocks == expected_max_input_blocks
+
+    max_output_blocks = max(
+        d["pipeline"].config.num_output_blocks[0] for _, d in rechunks
+    )
+    assert max_output_blocks == expected_max_output_blocks