Repository: dask/dask
File: dask/array/core.py (Python)

Code before:
            return_stored=return_stored,
            load_stored=load_stored,
        )
        layers[n] = map_layer
        if isinstance(t, Delayed):
            dependencies[n] = {sources_name, targets_name}
        else:
            dependencies[n] = {sources_name}
        map_keys += map_layer.keys()

    if return_stored:
        store_dsk = HighLevelGraph(layers, dependencies)
        load_store_dsk: HighLevelGraph | dict[tuple, Any] = store_dsk
        if compute:
            store_dlyds = [Delayed(k, store_dsk, layer=k[0]) for k in map_keys]
            store_dlyds = persist(*store_dlyds, **kwargs)
            store_dsk_2 = HighLevelGraph.merge(*[e.dask for e in store_dlyds])
            load_store_dsk = retrieve_from_ooc(map_keys, store_dsk, store_dsk_2)
            map_names = ["load-" + n for n in map_names]

        return tuple(
            Array(load_store_dsk, n, s.chunks, meta=s)
            for s, n in zip(sources, map_names)
        )

    elif compute:
        store_dsk = HighLevelGraph(layers, dependencies)
        compute_as_if_collection(Array, store_dsk, map_keys, **kwargs)
        return None

    else:
        key = "store-" + tokenize(map_names)
        layers[key] = {key: map_keys}
        dependencies[key] = set(map_names)
        store_dsk = HighLevelGraph(layers, dependencies)
        return Delayed(key, store_dsk)


def blockdims_from_blockshape(shape, chunks):
    """

    >>> blockdims_from_blockshape((10, 10), (4, 3))
    ((4, 4, 2), (3, 3, 3, 1))
    >>> blockdims_from_blockshape((10, 0), (4, 0))
    ((4, 4, 2), (0,))
    """
    if chunks is None:
        raise TypeError("Must supply chunks= keyword argument")
    if shape is None:
        raise TypeError("Must supply shape= keyword argument")
    if np.isnan(sum(shape)) or np.isnan(sum(chunks)):

Review comment (suggestion):
@jakirkham i was very confused by the default for `load_stored`. Why `not compute`? It would be useful to add a comment for the next person who comes around :)

Code after:
    sources_layer = Array.__dask_optimize__(
        sources_hlg, list(core.flatten([e.__dask_keys__() for e in sources]))
    )
    sources_name = "store-sources-" + tokenize(sources)
    layers = {sources_name: sources_layer}
    dependencies: dict[str, set[str]] = {sources_name: set()}

    # Optimize all targets together
    targets_keys = []
    targets_dsks = []
    for t in targets:
        if isinstance(t, Delayed):
            targets_keys.append(t.key)
            targets_dsks.append(t.__dask_graph__())
        elif is_dask_collection(t):
            raise TypeError("Targets must be either Delayed objects or array-likes")

    if targets_dsks:
        targets_hlg = HighLevelGraph.merge(*targets_dsks)
        targets_layer = Delayed.__dask_optimize__(targets_hlg, targets_keys)
        targets_name = "store-targets-" + tokenize(targets_keys)
        layers[targets_name] = targets_layer
        dependencies[targets_name] = set()

    if load_stored is None:
        load_stored = return_stored and not compute

    map_names = [
        "store-map-" + tokenize(s, t if isinstance(t, Delayed) else id(t), r)
        for s, t, r in zip(sources, targets, regions_list)
    ]
    map_keys: list[tuple] = []

    for s, t, n, r in zip(sources, targets, map_names, regions_list):
        map_layer = insert_to_ooc(
            keys=s.__dask_keys__(),
            chunks=s.chunks,
            out=t.key if isinstance(t, Delayed) else t,
            name=n,
            lock=lock,
            region=r,
            return_stored=return_stored,
            load_stored=load_stored,
        )
        layers[n] = map_layer
        if isinstance(t, Delayed):
            dependencies[n] = {sources_name, targets_name}
        else:
            dependencies[n] = {sources_name}
        map_keys += map_layer.keys()