Repository: tomwhite/cubed
Pull request #337 (2023-12-17 13:15:22)

Change the DAG to have separate nodes for operations and arrays
Currently there is a one-to-one correspondence between operations and arrays: one operation produces one array, and is represented by one node in the DAG.

However, this won't be true when we support multiple outputs (#69) where one operation can produce multiple arrays.

Similarly, it also won't support optimizations like sibling fusion (in [Apache Beam](https://beam.apache.org/releases/javadoc/2.0.0/org/apache/beam/sdk/transforms/ParDo.html), and from [FlumeJava](https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/35650.pdf) originally), where operations that have the same inputs and produce different outputs can be fused into one operation producing multiple outputs. ([Wukong](https://github.com/ds2-lab/Wukong) calls this "task clustering".)

So we need to change the DAG representation so that one operation can produce multiple arrays, and this means breaking  nodes into two types: operations and arrays.

This PR does exactly that. An example DAG (from `test_add_with_broadcast`) in the current representation

![cubed](https://github.com/tomwhite/cubed/assets/85085/40f78b62-38b0-4428-9399-170afc70d8c2)

now looks like this:

![cubed-new-dag](https://github.com/tomwhite/cubed/assets/85085/eb435aaf-0d60-42a4-ba94-e9293831119b)

Note:
* Operations are shown as boxes with rounded corners, arrays are rectangles.
* Colours indicate the following: pink for a blockwise operation, green for a rechunk operation, orange for a materialized array, white for an operation that is not run in its own pipeline or for an array that is not materialized to disk as a Zarr array ("virtual arrays")

I wondered about having an option to visualize DAGs as we currently do, but that won't work when there are multiple outputs, so it may be best just to have the more general representation. 


Patch:
diff --git a/cubed/array_api/creation_functions.py b/cubed/array_api/creation_functions.py
index 7aa6e0fb..c540bb95 100644
--- a/cubed/array_api/creation_functions.py
+++ b/cubed/array_api/creation_functions.py
@@ -7,9 +7,12 @@ from zarr.util import normalize_shape
 from cubed.backend_array_api import namespace as nxp
 from cubed.core import Plan, gensym
 from cubed.core.ops import map_direct
-from cubed.core.plan import new_temp_path
-from cubed.storage.virtual import virtual_empty, virtual_full, virtual_offsets
-from cubed.storage.zarr import lazy_from_array
+from cubed.storage.virtual import (
+    virtual_empty,
+    virtual_full,
+    virtual_in_memory,
+    virtual_offsets,
+)
 from cubed.utils import to_chunksize
 from cubed.vendor.dask.array.core import normalize_chunks
 
@@ -70,11 +73,9 @@ def asarray(
     if dtype is None:
         dtype = a.dtype
 
-    # write to zarr
     chunksize = to_chunksize(normalize_chunks(chunks, shape=a.shape, dtype=dtype))
     name = gensym()
-    zarr_path = new_temp_path(name=name, spec=spec)
-    target = lazy_from_array(a, dtype=dtype, chunks=chunksize, store=zarr_path)
+    target = virtual_in_memory(a, chunks=chunksize)
 
     plan = Plan._new(name, "asarray", target)
     return Array(name, target, spec, plan)
diff --git a/cubed/core/plan.py b/cubed/core/plan.py
index 2df1e185..f5999073 100644
--- a/cubed/core/plan.py
+++ b/cubed/core/plan.py
@@ -5,8 +5,8 @@ from datetime import datetime
 from functools import lru_cache
 
 import networkx as nx
+import zarr
 
-from cubed.backend_array_api import backend_array_to_numpy_array
 from cubed.primitive.blockwise import can_fuse_pipelines, fuse
 from cubed.runtime.pipeline import visit_nodes
 from cubed.runtime.types import CubedPipeline
@@ -16,6 +16,14 @@ from cubed.utils import chunk_memory, extract_stack_summaries, join_path, memory
 # A unique ID with sensible ordering, used for making directory names
 CONTEXT_ID = f"cubed-{datetime.now().strftime('%Y%m%dT%H%M%S')}-{uuid.uuid4()}"
 
+sym_counter = 0
+
+
+def gensym(name="op"):
+    global sym_counter
+    sym_counter += 1
+    return f"{name}-{sym_counter:03}"
+
 
 class Plan:
     """Deferred computation plan for a graph of arrays.
@@ -30,6 +38,12 @@ class Plan:
     a function with repeated inputs. For example, consider `equals` where the
     two arguments are the same array. We need to keep track of these cases, so
     we use a NetworkX `MultiDiGraph` rather than just as `DiGraph`.
+
+    Compared to a more traditional DAG representing a computation, in Cubed
+    nodes are not values that are passed to functions, they are instead
+    "parallel computations" which are run for their side effects. Data does
+    not flow through the graph - it is written to external storage (Zarr files)
+    as the output of one pipeline, then read back as the input to later pipelines.
     """
 
     def __init__(self, dag):
@@ -56,28 +70,50 @@ class Plan:
         frame = inspect.currentframe().f_back  # go back one in the stack
         stack_summaries = extract_stack_summaries(frame, limit=10)
 
+        op_name_unique = gensym()
+
         if pipeline is None:
+            # op
             dag.add_node(
-                name,
-                name=name,
+                op_name_unique,
+                name=op_name_unique,
                 op_name=op_name,
-                target=target,
+                type="op",
                 stack_summaries=stack_summaries,
                 hidden=hidden,
             )
-        else:
+            # array (when multiple outputs are supported there could be more than one)
             dag.add_node(
                 name,
                 name=name,
-                op_name=op_name,
+                type="array",
                 target=target,
+                hidden=hidden,
+            )
+            dag.add_edge(op_name_unique, name)
+        else:
+            # op
+            dag.add_node(
+                op_name_unique,
+                name=op_name_unique,
+                op_name=op_name,
+                type="op",
                 stack_summaries=stack_summaries,
                 hidden=hidden,
                 pipeline=pipeline,
             )
+            # array (when multiple outputs are supported there could be more than one)
+            dag.add_node(
+                name,
+                name=name,
+                type="array",
+                target=target,
+                hidden=hidden,
+            )
+            dag.add_edge(op_name_unique, name)
         for x in source_arrays:
             if hasattr(x, "name"):
-                dag.add_edge(x.name, name)
+                dag.add_edge(x.name, op_name_unique)
 
         return Plan(dag)
 
@@ -94,42 +130,62 @@ class Plan:
         nodes = {n: d for (n, d) in dag.nodes(data=True)}
 
         def can_fuse(n):
-            # node must have a single predecessor
-            #   - not multiple edges pointing to a single predecessor
-            # node must be the single successor to the predecessor
-            # and both must have pipelines that can be fused
-            if dag.in_degree(n) != 1:
+            # fuse a single chain looking like this:
+            # op1 -> op2_input -> op2
+
+            op2 = n
+
+            # if node (op2) does not have a pipeline then it can't be fused
+            if "pipeline" not in nodes[op2]:
+                return False
+
+            # if node (op2) does not have exactly one input then don't fuse
+            # (it could have no inputs or multiple inputs)
+            if dag.in_degree(op2) != 1:
+                return False
+
+            # if input is used by another node then don't fuse
+            op2_input = next(dag.predecessors(op2))
+            if dag.out_degree(op2_input) != 1:
                 return False
-            pre = next(dag.predecessors(n))
-            if dag.out_degree(pre) != 1:
+
+            # if node producing input (op1) has more than one output then don't fuse
+            op1 = next(dag.predecessors(op2_input))
+            if dag.out_degree(op1) != 1:
                 return False
-            if "pipeline" not in nodes[pre] or "pipeline" not in nodes[n]:
+
+            # op1 and op2 must have pipelines that can be fused
+            if "pipeline" not in nodes[op1]:
                 return False
-            return can_fuse_pipelines(nodes[pre]["pipeline"], nodes[n]["pipeline"])
+            return can_fuse_pipelines(nodes[op1]["pipeline"], nodes[op2]["pipeline"])
 
         for n in list(dag.nodes()):
             if can_fuse(n):
-                pre = next(dag.predecessors(n))
-                pipeline = fuse(nodes[pre]["pipeline"], nodes[n]["pipeline"])
-                nodes[n]["pipeline"] = pipeline
-                assert nodes[n]["target"] == pipeline.target_array
+                op2 = n
+                op2_input = next(dag.predecessors(op2))
+                op1 = next(dag.predecessors(op2_input))
+                op1_inputs = list(dag.predecessors(op1))
+
+                pipeline = fuse(nodes[op1]["pipeline"], nodes[op2]["pipeline"])
+                nodes[op2]["pipeline"] = pipeline
 
-                for p in dag.predecessors(pre):
-                    dag.add_edge(p, n)
-                dag.remove_node(pre)
+                for n in op1_inputs:
+                    dag.add_edge(n, op2)
+                dag.remove_node(op2_input)
+                dag.remove_node(op1)
 
         return Plan(dag)
 
     def _create_lazy_zarr_arrays(self, dag):
         # find all lazy zarr arrays in dag
-        all_array_nodes = []
+        all_pipeline_nodes = []
         lazy_zarr_arrays = []
         reserved_mem_values = []
         for n, d in dag.nodes(data=True):
             if "pipeline" in d and d["pipeline"].reserved_mem is not None:
                 reserved_mem_values.append(d["pipeline"].reserved_mem)
-            if isinstance(d["target"], LazyZarrArray):
-                all_array_nodes.append(n)
+                all_pipeline_nodes.append(n)
+            if "target" in d and isinstance(d["target"], LazyZarrArray):
                 lazy_zarr_arrays.append(d["target"])
 
         reserved_mem = max(reserved_mem_values, default=0)
@@ -143,14 +199,20 @@ class Plan:
                 name,
                 name=name,
                 op_name=op_name,
-                target=None,
+                type="op",
                 pipeline=pipeline,
                 projected_mem=pipeline.projected_mem,
                 num_tasks=pipeline.num_tasks,
             )
-            # make create arrays node a dependency of all lazy array nodes
-            for n in all_array_nodes:
-                dag.add_edge(name, n)
+            dag.add_node(
+                "arrays",
+                name="arrays",
+                target=None,
+            )
+            dag.add_edge(name, "arrays")
+            # make create arrays node a predecessor of all pipeline nodes so it runs first
+            for n in all_pipeline_nodes:
+                dag.add_edge("arrays", n)
 
         return dag
 
@@ -197,7 +259,7 @@ class Plan:
     def num_arrays(self, optimize_graph: bool = True) -> int:
         """Return the number of arrays in this plan."""
         dag = self._finalize_dag(optimize_graph=optimize_graph)
-        return sum(n != "create-arrays" for n in dag.nodes())
+        return sum(d.get("type") == "array" for _, d in dag.nodes(data=True))
 
     def max_projected_mem(self, optimize_graph=True, resume=None):
         """Return the maximum projected memory across all tasks to execute this plan."""
@@ -214,8 +276,8 @@ class Plan:
         dag = self._finalize_dag(optimize_graph=optimize_graph)
         dag = dag.copy()  # make a copy since we mutate the DAG below
 
-        # remove edges from create-arrays node to avoid cluttering the diagram
-        dag.remove_edges_from(list(dag.out_edges("create-arrays")))
+        # remove edges from create-arrays output node to avoid cluttering the diagram
+        dag.remove_edges_from(list(dag.out_edges("arrays")))
 
         # remove hidden nodes
         dag.remove_nodes_from(
@@ -253,69 +315,85 @@ class Plan:
 
         # now set node attributes with visualization info
         for n, d in dag.nodes(data=True):
-            if d["op_name"] == "blockwise":
-                d["style"] = "filled"
-                d["fillcolor"] = "#dcbeff"
-                op_name_summary = "(bw)"
-            elif d["op_name"] == "rechunk":
-                d["style"] = "filled"
-                d["fillcolor"] = "#aaffc3"
-                op_name_summary = "(rc)"
-            else:  # creation function
-                op_name_summary = ""
-            target = d["target"]
-            if target is not None:
+            tooltip = f"name: {n}\n"
+            node_type = d.get("type", None)
+            if node_type == "op":
+                op_name = d["op_name"]
+                if op_name == "blockwise":
+                    d["style"] = '"rounded,filled"'
+                    d["fillcolor"] = "#dcbeff"
+                    op_name_summary = "(bw)"
+                elif op_name == "rechunk":
+                    d["style"] = '"rounded,filled"'
+                    d["fillcolor"] = "#aaffc3"
+                    op_name_summary = "(rc)"
+                else:
+                    # creation function
+                    d["style"] = "rounded"
+                    op_name_summary = ""
+                tooltip += f"op: {op_name}"
+
+                if "pipeline" in d:
+                    pipeline = d["pipeline"]
+                    tooltip += (
+                        f"\nprojected memory: {memory_repr(pipeline.projected_mem)}"
+                    )
+                    tooltip += f"\ntasks: {pipeline.num_tasks}"
+                    if pipeline.write_chunks is not None:
+                        tooltip += f"\nwrite chunks: {pipeline.write_chunks}"
+
+                    # remove pipeline attribute since it is a long string that causes graphviz to fail
+                    del d["pipeline"]
+
+                if "stack_summaries" in d and d["stack_summaries"] is not None:
+                    # add call stack information
+                    stack_summaries = d["stack_summaries"]
+
+                    first_cubed_i = min(
+                        i for i, s in enumerate(stack_summaries) if s.is_cubed()
+                    )
+                    first_cubed_summary = stack_summaries[first_cubed_i]
+                    caller_summary = stack_summaries[first_cubed_i - 1]
+
+                    d["label"] = f"{first_cubed_summary.name} {op_name_summary}"
+
+                    calls = " -> ".join(
+                        [
+                            s.name
+                            for s in stack_summaries
+                            if not s.is_on_python_lib_path()
+                        ]
+                    )
+
+                    line = f"{caller_summary.lineno} in {caller_summary.name}"
+
+                    tooltip += f"\ncalls: {calls}"
+                    tooltip += f"\nline: {line}"
+                    del d["stack_summaries"]
+
+            elif node_type == "array":
+                target = d["target"]
                 chunkmem = memory_repr(chunk_memory(target.dtype, target.chunks))
-                tooltip = (
-                    f"name: {n}\n"
-                    f"shape: {target.shape}\n"
-                    f"chunks: {target.chunks}\n"
-                    f"dtype: {target.dtype}\n"
-                    f"chunk memory: {chunkmem}\n"
-                )
-            else:
-                tooltip = ""
-            if "pipeline" in d:
-                pipeline = d["pipeline"]
-                tooltip += f"\nprojected memory: {memory_repr(pipeline.projected_mem)}"
-                tooltip += f"\ntasks: {pipeline.num_tasks}"
-                if pipeline.write_chunks is not None:
-                    tooltip += f"\nwrite chunks: {pipeline.write_chunks}"
-            if "stack_summaries" in d and d["stack_summaries"] is not None:
-                # add call stack information
-                stack_summaries = d["stack_summaries"]
-
-                first_cubed_i = min(
-                    i for i, s in enumerate(stack_summaries) if s.is_cubed()
-                )
-                first_cubed_summary = stack_summaries[first_cubed_i]
-                caller_summary = stack_summaries[first_cubed_i - 1]
 
+                # materialized arrays are light orange, virtual arrays are white
+                if isinstance(target, (LazyZarrArray, zarr.Array)):
+                    d["style"] = "filled"
+                    d["fillcolor"] = "#ffd8b1"
                 if n in array_display_names:
-                    var_name = f" ({array_display_names[n]})"
+                    var_name = array_display_names[n]
+                    d["label"] = f"{n} ({var_name})"
+                    tooltip += f"variable: {var_name}\n"
                 else:
-                    var_name = ""
-                d[
-                    "label"
-                ] = f"{n}{var_name}\n{first_cubed_summary.name} {op_name_summary}"
+                    d["label"] = n
+                tooltip += f"shape: {target.shape}\n"
+                tooltip += f"chunks: {target.chunks}\n"
+                tooltip += f"dtype: {target.dtype}\n"
+                tooltip += f"chunk memory: {chunkmem}\n"
 
-                calls = " -> ".join(
-                    [s.name for s in stack_summaries if not s.is_on_python_lib_path()]
-                )
-
-                line = f"{caller_summary.lineno} in {caller_summary.name}"
-
-                tooltip += f"\ncalls: {calls}"
-                tooltip += f"\nline: {line}"
-                del d["stack_summaries"]
+                del d["target"]
 
             d["tooltip"] = tooltip.strip()
 
-            # remove pipeline attribute since it is a long string that causes graphviz to fail
-            if "pipeline" in d:
-                del d["pipeline"]
-            if "target" in d:
-                del d["target"]
             if "name" in d:  # pydot already has name
                 del d["name"]
         gv = nx.drawing.nx_pydot.to_pydot(dag)
@@ -364,19 +442,9 @@ def create_zarr_array(lazy_zarr_array, *, config=None):
 
 
 def create_zarr_arrays(lazy_zarr_arrays, reserved_mem):
-    # projected memory is size of largest initial values, or dtype size if there aren't any
+    # projected memory is size of largest dtype size (for a fill value)
     projected_mem = (
-        max(
-            [
-                # TODO: calculate nbytes from size and dtype itemsize
-                backend_array_to_numpy_array(lza.initial_values).nbytes
-                if lza.initial_values is not None
-                else lza.dtype.itemsize
-                for lza in lazy_zarr_arrays
-            ],
-            default=0,
-        )
-        + reserved_mem
+        max([lza.dtype.itemsize for lza in lazy_zarr_arrays], default=0) + reserved_mem
     )
     num_tasks = len(lazy_zarr_arrays)
 
diff --git a/cubed/runtime/executors/beam.py b/cubed/runtime/executors/beam.py
index fbaf4cd7..fd919c0c 100644
--- a/cubed/runtime/executors/beam.py
+++ b/cubed/runtime/executors/beam.py
@@ -89,9 +89,10 @@ class BeamDagExecutor(DagExecutor):
         pipeline = beam.Pipeline(**kwargs)
 
         for name, node in visit_nodes(dag, resume=resume):
-            rechunker_pipeline = node["pipeline"]
+            cubed_pipeline = node["pipeline"]
 
-            dep_nodes = list(dag.predecessors(name))
+            inputs = list(dag.predecessors(name))
+            dep_nodes = list(p for n in inputs for p in dag.predecessors(n))
 
             pcolls = [
                 p
@@ -100,17 +101,17 @@ class BeamDagExecutor(DagExecutor):
             ]
             if len(pcolls) == 0:
                 pcoll = pipeline | gensym("Start") >> beam.Create([-1])
-                pcoll = add_to_pcoll(name, rechunker_pipeline, pcoll)
+                pcoll = add_to_pcoll(name, cubed_pipeline, pcoll)
                 dag.nodes[name]["pcoll"] = pcoll
 
             elif len(pcolls) == 1:
                 pcoll = pcolls[0]
-                pcoll = add_to_pcoll(name, rechunker_pipeline, pcoll)
+                pcoll = add_to_pcoll(name, cubed_pipeline, pcoll)
                 dag.nodes[name]["pcoll"] = pcoll
             else:
                 pcoll = pcolls | gensym("Flatten") >> beam.Flatten()
                 pcoll |= gensym("Distinct") >> beam.Distinct()
-                pcoll = add_to_pcoll(name, rechunker_pipeline, pcoll)
+                pcoll = add_to_pcoll(name, cubed_pipeline, pcoll)
                 dag.nodes[name]["pcoll"] = pcoll
 
         result = pipeline.run()
@@ -121,23 +122,23 @@ class BeamDagExecutor(DagExecutor):
             wait_until_finish_with_callbacks(result, callbacks)
 
 
-def add_to_pcoll(name, rechunker_pipeline, pcoll):
+def add_to_pcoll(name, cubed_pipeline, pcoll):
     step = 0
     stage = Stage(
-        rechunker_pipeline.function,
-        rechunker_pipeline.name,
-        rechunker_pipeline.mappable,
+        cubed_pipeline.function,
+        cubed_pipeline.name,
+        cubed_pipeline.mappable,
     )
     if stage.mappable is not None:
         pcoll |= stage.name >> _SingleArgumentStage(
-            step, stage, rechunker_pipeline.config, name
+            step, stage, cubed_pipeline.config, name
         )
     else:
         pcoll |= stage.name >> beam.Map(
             _no_arg_stage,
             current=step,
             fun=stage.function,
-            config=rechunker_pipeline.config,
+            config=cubed_pipeline.config,
         )
 
     # This prevents fusion:
diff --git a/cubed/runtime/pipeline.py b/cubed/runtime/pipeline.py
index eaf3aaeb..0b39bd9a 100644
--- a/cubed/runtime/pipeline.py
+++ b/cubed/runtime/pipeline.py
@@ -5,21 +5,32 @@ import networkx as nx
 from cubed.storage.zarr import open_if_lazy_zarr_array
 
 
-def already_computed(node_dict: Dict[str, Any], resume: Optional[bool] = None) -> bool:
+def already_computed(
+    name, dag, nodes: Dict[str, Any], resume: Optional[bool] = None
+) -> bool:
     """
     Return True if the array for a node doesn't have a pipeline to compute it,
-    or it has already been computed (all chunks are present).
+    or if all its outputs have already been computed (all chunks are present).
     """
-    pipeline = node_dict.get("pipeline", None)
+    pipeline = nodes[name].get("pipeline", None)
     if pipeline is None:
         return True
 
-    target = node_dict.get("target", None)
-    if resume and target is not None:
-        target = open_if_lazy_zarr_array(target)
-        # this check can be expensive since it has to list the directory to find nchunks_initialized
-        if target.ndim > 0 and target.nchunks_initialized == target.nchunks:
-            return True
+    # if no outputs have targets then need to compute (this is the create-arrays case)
+    if all(
+        [nodes[output].get("target", None) is None for output in dag.successors(name)]
+    ):
+        return False
+
+    if resume:
+        for output in dag.successors(name):
+            target = nodes[output].get("target", None)
+            if target is not None:
+                target = open_if_lazy_zarr_array(target)
+                # this check can be expensive since it has to list the directory to find nchunks_initialized
+                if target.ndim == 0 or target.nchunks_initialized != target.nchunks:
+                    return False
+        return True
 
     return False
 
@@ -28,7 +39,7 @@ def visit_nodes(dag, resume=None):
     """Return a generator that visits the nodes in the DAG in topological order."""
     nodes = {n: d for (n, d) in dag.nodes(data=True)}
     for name in list(nx.topological_sort(dag)):
-        if already_computed(nodes[name], resume=resume):
+        if already_computed(name, dag, nodes, resume=resume):
             continue
         yield name, nodes[name]
 
@@ -40,7 +51,7 @@ def visit_node_generations(dag, resume=None):
         gen = [
             (name, nodes[name])
             for name in names
-            if not already_computed(nodes[name], resume=resume)
+            if not already_computed(name, dag, nodes, resume=resume)
         ]
         if len(gen) > 0:
             yield gen
diff --git a/cubed/storage/virtual.py b/cubed/storage/virtual.py
index 2a48f6ec..2b3bf489 100644
--- a/cubed/storage/virtual.py
+++ b/cubed/storage/virtual.py
@@ -8,6 +8,7 @@ from zarr.indexing import BasicIndexer, is_slice
 from cubed.backend_array_api import namespace as nxp
 from cubed.backend_array_api import numpy_array_to_backend_array
 from cubed.types import T_DType, T_RegularChunks, T_Shape
+from cubed.utils import memory_repr
 
 
 class VirtualEmptyArray:
@@ -97,6 +98,43 @@ class VirtualOffsetsArray:
         )
 
 
+class VirtualInMemoryArray:
+    """A small array that is held in memory but never materialized on disk."""
+
+    def __init__(
+        self,
+        array: np.ndarray,  # TODO: generalise
+        chunks: T_RegularChunks,
+        max_nbytes: int = 10**6,
+    ):
+        if array.nbytes > max_nbytes:
+            raise ValueError(
+                f"Size of in memory array is {memory_repr(array.nbytes)} which exceeds maximum of {memory_repr(max_nbytes)}. Consider loading the array from storage using `from_array`."
+            )
+        self.array = array
+        # use an in-memory Zarr array as a template since it normalizes its properties
+        # and is needed for oindex
+        template = zarr.empty(
+            array.shape,
+            dtype=array.dtype,
+            chunks=chunks,
+            store=zarr.storage.MemoryStore(),
+        )
+        self.shape = template.shape
+        self.dtype = template.dtype
+        self.chunks = template.chunks
+        self.template = template
+        if array.size > 0:
+            template[...] = array
+
+    def __getitem__(self, key):
+        return self.array.__getitem__(key)
+
+    @property
+    def oindex(self):
+        return self.template.oindex
+
+
 def _key_to_index_tuple(selection):
     if isinstance(selection, slice):
         selection = (selection,)
@@ -131,3 +169,10 @@ def virtual_full(
 
 def virtual_offsets(shape: T_Shape) -> VirtualOffsetsArray:
     return VirtualOffsetsArray(shape)
+
+
+def virtual_in_memory(
+    array: np.ndarray,
+    chunks: T_RegularChunks,
+) -> VirtualInMemoryArray:
+    return VirtualInMemoryArray(array, chunks)
diff --git a/cubed/storage/zarr.py b/cubed/storage/zarr.py
index 60c33874..d1f17e05 100644
--- a/cubed/storage/zarr.py
+++ b/cubed/storage/zarr.py
@@ -1,9 +1,7 @@
-from typing import Any, Optional, Union
+from typing import Any, Union
 
 import zarr
-from numpy import ndarray
 
-from cubed.backend_array_api import backend_array_to_numpy_array
 from cubed.types import T_DType, T_RegularChunks, T_Shape, T_Store
 
 
@@ -21,7 +19,6 @@ class LazyZarrArray:
         dtype: T_DType,
         chunks: T_RegularChunks,
         store: T_Store,
-        initial_values: Optional[ndarray] = None,
         fill_value: Any = None,
         **kwargs,
     ):
@@ -35,7 +32,6 @@ class LazyZarrArray:
         self.chunks = template.chunks
 
         self.store = store
-        self.initial_values = initial_values
         self.fill_value = fill_value
         self.kwargs = kwargs
 
@@ -60,8 +56,6 @@ class LazyZarrArray:
             fill_value=self.fill_value,
             **self.kwargs,
         )
-        if self.initial_values is not None and self.initial_values.size > 0:
-            target[...] = backend_array_to_numpy_array(self.initial_values)
         return target
 
     def open(self) -> zarr.Array:
@@ -91,14 +85,6 @@ def lazy_empty(
     return LazyZarrArray(shape, dtype, chunks, store, **kwargs)
 
 
-def lazy_from_array(
-    array: ndarray, *, dtype: T_DType, chunks: T_RegularChunks, store: T_Store, **kwargs
-) -> LazyZarrArray:
-    return LazyZarrArray(
-        array.shape, dtype, chunks, store, initial_values=array, **kwargs
-    )
-
-
 def lazy_full(
     shape: T_Shape,
     fill_value: Any,


Test patch:
diff --git a/cubed/tests/storage/test_zarr.py b/cubed/tests/storage/test_zarr.py
index c7da964b..e4a6de44 100644
--- a/cubed/tests/storage/test_zarr.py
+++ b/cubed/tests/storage/test_zarr.py
@@ -2,7 +2,7 @@ import numpy as np
 import pytest
 from numpy.testing import assert_array_equal
 
-from cubed.storage.zarr import lazy_empty, lazy_from_array, lazy_full
+from cubed.storage.zarr import lazy_empty, lazy_full
 
 
 def test_lazy_empty(tmp_path):
@@ -18,20 +18,6 @@ def test_lazy_empty(tmp_path):
     arr.open()
 
 
-def test_lazy_from_array(tmp_path):
-    zarr_path = tmp_path / "lazy.zarr"
-    a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=int)
-    arr = lazy_from_array(a, dtype=a.dtype, chunks=(2, 2), store=zarr_path)
-
-    assert not zarr_path.exists()
-    with pytest.raises(ValueError):
-        arr.open()
-
-    arr.create()
-    assert zarr_path.exists()
-    assert_array_equal(arr.open()[:], a)
-
-
 def test_lazy_full(tmp_path):
     zarr_path = tmp_path / "lazy.zarr"
     arr = lazy_full(
diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py
index 85e3a3f5..63902870 100644
--- a/cubed/tests/test_core.py
+++ b/cubed/tests/test_core.py
@@ -40,6 +40,15 @@ def modal_executor(request):
     return request.param
 
 
+def test_as_array_fails(spec):
+    a = np.ones((1000, 1000))
+    with pytest.raises(
+        ValueError,
+        match="Size of in memory array is 8.0 MB which exceeds maximum of 1.0 MB.",
+    ):
+        xp.asarray(a, chunks=(100, 100), spec=spec)
+
+
 def test_regular_chunks(spec):
     xp.ones((5, 5), chunks=((2, 2, 1), (5,)), spec=spec)
     with pytest.raises(ValueError):
@@ -220,9 +229,8 @@ def test_rechunk_same_chunks(spec):
     b = a.rechunk((2, 1))
     task_counter = TaskCounter()
     res = b.compute(callbacks=[task_counter])
-    # no tasks except array creation task should have run since chunks are same
-    num_created_arrays = 1
-    assert task_counter.value == num_created_arrays
+    # no tasks should have run since chunks are same
+    assert task_counter.value == 0
 
     assert_array_equal(res, np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
 
@@ -323,7 +331,9 @@ def test_reduction_multiple_rounds(tmp_path, executor):
     b = xp.sum(a, axis=0, dtype=np.uint8)
     # check that there is > 1 blockwise step (after optimization)
     blockwises = [
-        n for (n, d) in b.plan.dag.nodes(data=True) if d["op_name"] == "blockwise"
+        n
+        for (n, d) in b.plan.dag.nodes(data=True)
+        if d.get("op_name", None) == "blockwise"
     ]
     assert len(blockwises) > 1
     assert b.plan.max_projected_mem() <= 1000
diff --git a/cubed/tests/test_executor_features.py b/cubed/tests/test_executor_features.py
index 944640d4..9117278a 100644
--- a/cubed/tests/test_executor_features.py
+++ b/cubed/tests/test_executor_features.py
@@ -93,7 +93,7 @@ def test_callbacks(spec, executor):
         np.array([[2, 3, 4], [5, 6, 7], [8, 9, 10]]),
     )
 
-    num_created_arrays = 3
+    num_created_arrays = 1
     assert task_counter.value == num_created_arrays + 4
 
 
@@ -111,7 +111,7 @@ def test_callbacks_modal(spec, modal_executor):
             np.array([[2, 3, 4], [5, 6, 7], [8, 9, 10]]),
         )
 
-        num_created_arrays = 3
+        num_created_arrays = 1
         assert task_counter.value == num_created_arrays + 4
     finally:
         fs = fsspec.open(tmp_path).fs
@@ -132,12 +132,12 @@ def test_resume(spec, executor):
     c = xp.add(a, b)
     d = xp.negative(c)
 
-    num_created_arrays = 4  # a, b, c, d
+    num_created_arrays = 2  # c, d
     assert d.plan.num_tasks(optimize_graph=False) == num_created_arrays + 8
 
     task_counter = TaskCounter()
     c.compute(executor=executor, callbacks=[task_counter], optimize_graph=False)
-    num_created_arrays = 3  # a, b, c
+    num_created_arrays = 1  # c
     assert task_counter.value == num_created_arrays + 4
 
     # since c has already been computed, when computing d only 4 tasks are run, instead of 8
@@ -146,7 +146,7 @@ def test_resume(spec, executor):
         executor=executor, callbacks=[task_counter], optimize_graph=False, resume=True
     )
     # the create arrays tasks are run again, even though they exist
-    num_created_arrays = 4  # a, b, c, d
+    num_created_arrays = 2  # c, d
     assert task_counter.value == num_created_arrays + 4
 
 
diff --git a/cubed/tests/test_optimization.py b/cubed/tests/test_optimization.py
index 7a538ce6..803a8777 100644
--- a/cubed/tests/test_optimization.py
+++ b/cubed/tests/test_optimization.py
@@ -18,11 +18,13 @@ def test_fusion(spec):
     c = xp.astype(b, np.float32)
     d = xp.negative(c)
 
-    num_created_arrays = 4  # a, b, c, d
-    assert d.plan.num_arrays(optimize_graph=False) == num_created_arrays
+    num_arrays = 4  # a, b, c, d
+    num_created_arrays = 3  # b, c, d (a is not created on disk)
+    assert d.plan.num_arrays(optimize_graph=False) == num_arrays
     assert d.plan.num_tasks(optimize_graph=False) == num_created_arrays + 12
-    num_created_arrays = 2  # a, d
-    assert d.plan.num_arrays(optimize_graph=True) == num_created_arrays
+    num_arrays = 2  # a, d
+    num_created_arrays = 1  # d (a is not created on disk)
+    assert d.plan.num_arrays(optimize_graph=True) == num_arrays
     assert d.plan.num_tasks(optimize_graph=True) == num_created_arrays + 4
 
     task_counter = TaskCounter()
@@ -41,9 +43,9 @@ def test_fusion_transpose(spec):
     c = xp.astype(b, np.float32)
     d = c.T
 
-    num_created_arrays = 4  # a, b, c, d
+    num_created_arrays = 3  # b, c, d
     assert d.plan.num_tasks(optimize_graph=False) == num_created_arrays + 12
-    num_created_arrays = 2  # a, d
+    num_created_arrays = 1  # d
     assert d.plan.num_tasks(optimize_graph=True) == num_created_arrays + 4
 
     task_counter = TaskCounter()