Change list of output arrays in BlockwiseSpec to a map (#724)

Author: tomwhite

cubed/core/ops.py

@@ -306,6 +306,7 @@ def blockwise(
         reserved_mem=spec.reserved_mem,
         extra_projected_mem=extra_projected_mem,
         target_store=target_store,
+        target_name=name,
         target_path=target_path,
         storage_options=spec.storage_options,
         compressor=spec.zarr_compressor,
@@ -314,7 +315,6 @@ def blockwise(
         chunks=_chunks,
         new_axes=new_axes,
         in_names=in_names,
-        out_name=name,
         buffer_copies=buffer_copies,
         extra_func_kwargs=extra_func_kwargs,
         fusable_with_predecessors=fusable_with_predecessors,
@@ -454,10 +454,12 @@ def _general_blockwise(
             ts if ts is not None else new_temp_path(name=n, spec=spec)
             for n, ts in zip(name, target_stores)
         ]
+        target_names = name
     else:  # single output
         name = gensym()
         if target_stores is None:
             target_stores = [new_temp_path(name=name, spec=spec)]
+        target_names = [name]
 
     op = primitive_general_blockwise(
         func,
@@ -468,6 +470,7 @@ def _general_blockwise(
         extra_projected_mem=extra_projected_mem,
         buffer_copies=buffer_copies,
         target_stores=target_stores,
+        target_names=target_names,
         target_paths=target_paths,
         storage_options=spec.storage_options,
         compressor=spec.zarr_compressor,

cubed/primitive/blockwise.py

@@ -63,8 +63,9 @@ class BlockwiseSpec:
         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]
-        Write proxy list where entries have an ``array`` attribute that supports ``__setitem__``.
+    writes_map : Dict[str, CubedArrayProxy]
+        Write proxy dictionary keyed by array name where entries have an ``array`` attribute that supports ``__setitem__``.
+        Note that the order of the entries must correspond to the order of the outputs created by the function.
     """
 
     key_function: Callable[..., Any]
@@ -74,7 +75,7 @@ class BlockwiseSpec:
     num_output_blocks: Tuple[int, ...]
     iterable_input_blocks: Tuple[bool, ...]
     reads_map: Dict[str, CubedArrayProxy]
-    writes_list: List[CubedArrayProxy]
+    writes_map: Dict[str, CubedArrayProxy]
     return_writes_stores: bool = False
 
 
@@ -92,22 +93,20 @@ def apply_blockwise(
         results, tuple
     ):
         results = (results,)
-    for i, result in enumerate(results):
+    for result, write_proxy in zip(results, config.writes_map.values()):
         out_chunk_key = key_to_slices(
-            out_coords_tuple, config.writes_list[i].array, config.writes_list[i].chunks
+            out_coords_tuple, write_proxy.array, write_proxy.chunks
         )
         if isinstance(result, dict):  # group of arrays with named fields
             for k, v in result.items():
                 v = backend_array_to_numpy_array(v)
-                config.writes_list[i].open().set_basic_selection(
-                    out_chunk_key, v, fields=k
-                )
+                write_proxy.open().set_basic_selection(out_chunk_key, v, fields=k)
         else:
             result = backend_array_to_numpy_array(result)
-            config.writes_list[i].open()[out_chunk_key] = result
+            write_proxy.open()[out_chunk_key] = result
 
     if config.return_writes_stores:
-        return [write_proxy.open().store for write_proxy in config.writes_list]
+        return [write_proxy.open().store for write_proxy in config.writes_map.values()]
     return None
 
 
@@ -152,6 +151,7 @@ def blockwise(
     allowed_mem: int,
     reserved_mem: int,
     target_store: T_Store,
+    target_name: str,
     target_path: Optional[str] = None,
     storage_options: Optional[Dict[str, Any]] = None,
     compressor: Union[dict, str, None] = "default",
@@ -160,7 +160,6 @@ def blockwise(
     chunks: T_Chunks,
     new_axes: Optional[Dict[int, int]] = None,
     in_names: Optional[List[str]] = None,
-    out_name: Optional[str] = None,
     extra_projected_mem: int = 0,
     buffer_copies: Optional[BufferCopies] = None,
     extra_func_kwargs: Optional[Dict[str, Any]] = None,
@@ -231,7 +230,7 @@ def blockwise(
 
     key_function = make_blockwise_key_function_flattened(
         func,
-        out_name or "out",
+        target_name,
         out_ind,
         *argindsstr,
         numblocks=numblocks,
@@ -245,6 +244,7 @@ def blockwise(
         allowed_mem=allowed_mem,
         reserved_mem=reserved_mem,
         target_stores=[target_store],
+        target_names=[target_name],
         target_paths=[target_path] if target_path is not None else None,
         storage_options=storage_options,
         compressor=compressor,
@@ -270,6 +270,7 @@ def general_blockwise(
     allowed_mem: int,
     reserved_mem: int,
     target_stores: List[T_Store],
+    target_names: List[str],
     target_paths: Optional[List[str]] = None,
     storage_options: Optional[Dict[str, Any]] = None,
     compressor: Union[dict, str, None] = "default",
@@ -341,7 +342,7 @@ def general_blockwise(
         name: CubedArrayProxy(array, array.chunks) for name, array in array_map.items()
     }
 
-    write_proxies = []
+    write_proxies = {}
     output_chunk_memory = 0
     target_arrays = []
 
@@ -372,7 +373,7 @@ def general_blockwise(
             )
         target_arrays.append(ta)
 
-        write_proxies.append(CubedArrayProxy(ta, chunksize))
+        write_proxies[target_names[i]] = CubedArrayProxy(ta, chunksize)
 
         # only one output chunk is read into memory at a time, so we find the largest
         output_chunk_memory = max(
@@ -559,7 +560,7 @@ def fused_func(*args):
 
     function_nargs = pipeline1.config.function_nargs
     read_proxies = pipeline1.config.reads_map
-    write_proxies = pipeline2.config.writes_list
+    write_proxies = pipeline2.config.writes_map
     return_writes_stores = pipeline2.config.return_writes_stores
     num_input_blocks = tuple(
         n * pipeline2.config.num_input_blocks[0]
@@ -621,7 +622,7 @@ def fuse_multiple(
         num_output_blocks=(1,),
         iterable_input_blocks=(False,),
         reads_map={},
-        writes_list=[],
+        writes_map={},
     )
     predecessor_bw_specs = [
         primitive_op.pipeline.config
@@ -710,7 +711,7 @@ def fuse_blockwise_specs(
     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
+    write_proxies = bw_spec.writes_map
     return_writes_stores = bw_spec.return_writes_stores
     return BlockwiseSpec(
         fused_key_func,

cubed/tests/primitive/test_blockwise.py

@@ -40,6 +40,7 @@ def test_blockwise(tmp_path, executor, reserved_mem):
         allowed_mem=allowed_mem,
         reserved_mem=reserved_mem,
         target_store=target_store,
+        target_name="target",
         shape=(3, 3),
         dtype=int,
         chunks=(2, 2),
@@ -70,7 +71,7 @@ def test_blockwise(tmp_path, executor, reserved_mem):
     assert_array_equal(res[:], np.outer([0, 1, 2], [10, 50, 100]))
 
 
-def _permute_dims(x, /, axes, allowed_mem, reserved_mem, target_store):
+def _permute_dims(x, /, axes, allowed_mem, reserved_mem, target_store, target_name):
     # From dask transpose
     if axes:
         if len(axes) != x.ndim:
@@ -86,6 +87,7 @@ def _permute_dims(x, /, axes, allowed_mem, reserved_mem, target_store):
         allowed_mem=allowed_mem,
         reserved_mem=reserved_mem,
         target_store=target_store,
+        target_name=target_name,
         shape=x.shape,
         dtype=x.dtype,
         chunks=x.chunks,
@@ -109,6 +111,7 @@ def test_blockwise_with_args(tmp_path, executor):
         allowed_mem=allowed_mem,
         reserved_mem=0,
         target_store=target_store,
+        target_name="target",
     )
 
     assert op.target_array.shape == (3, 3)
@@ -160,6 +163,7 @@ def test_blockwise_allowed_mem_exceeded(tmp_path, reserved_mem):
             allowed_mem=allowed_mem,
             reserved_mem=reserved_mem,
             target_store=target_store,
+            target_name="target",
             shape=(3, 3),
             dtype=np.int64,
             chunks=(2, 2),
@@ -204,6 +208,7 @@ def key_function(out_key):
         allowed_mem=allowed_mem,
         reserved_mem=0,
         target_stores=[target_store],
+        target_names=["target"],
         shapes=[(20,)],
         dtypes=[int],
         chunkss=[(6,)],
@@ -252,6 +257,7 @@ def block_function(out_key):
         allowed_mem=allowed_mem,
         reserved_mem=0,
         target_stores=[target_store1, target_store2],
+        target_names=["target1", "target2"],
         shapes=[(3, 3), (3, 3)],
         dtypes=[float, float],
         chunkss=[(2, 2), (2, 2)],
@@ -317,6 +323,7 @@ def block_function(out_key):
             allowed_mem=allowed_mem,
             reserved_mem=0,
             target_stores=[target_store1, target_store2],
+            target_names=["target1", "target2"],
             shapes=[(3, 3), (3, 3)],
             dtypes=[float, float],
             chunkss=[(2, 2), (4, 2)],  # numblocks differ

cubed/tests/primitive/test_blockwise_fusion.py

@@ -121,7 +121,7 @@ def make_blockwise_spec(
         num_output_blocks=num_output_blocks,
         iterable_input_blocks=iterable_input_blocks,
         reads_map={},  # unused
-        writes_list=[],  # unused
+        writes_map={},  # unused
     )