Function _make_blockwise_graph

dask/blockwise.py:914–996

(
    task,
    output,
    out_indices,
    *arrind_pairs,
    numblocks=None,
    concatenate=None,
    new_axes=None,
    output_blocks=None,
    dims=None,
    io_deps=None,
    keys=None,
)

def _make_blockwise_graph(
    task,
    output,
    out_indices,
    *arrind_pairs,
    numblocks=None,
    concatenate=None,
    new_axes=None,
    output_blocks=None,
    dims=None,
    io_deps=None,
    keys=None,
):
    if numblocks is None:
        raise ValueError("Missing required numblocks argument.")
    new_axes = new_axes or {}
    io_deps = io_deps or {}
    argpairs = list(toolz.partition(2, arrind_pairs))

    if concatenate is True:
        from dask.array.core import concatenate_axes as concatenate

    # Dictionary mapping {i: 3, j: 4, ...} for i, j, ... the dimensions
    dims = dims or _make_dims(argpairs, numblocks, new_axes)

    # Generate the abstract "plan" before constructing
    # the actual graph
    coord_maps, concat_axes, dummies = _get_coord_mapping(
        dims,
        out_indices,
        numblocks,
        argpairs,
        concatenate,
    )

    # Apply Culling.
    # Only need to construct the specified set of output blocks.
    # Note that we must convert itertools.product to list,
    # because we may need to loop through output_blocks more than
    # once below (itertools.product already uses an internal list,
    # so this is not a memory regression)
    output_blocks = output_blocks or list(
        itertools.product(*[range(dims[i]) for i in out_indices])
    )
    from dask._task_spec import DataNode

    dsk = {}
    for out_coords in output_blocks:
        this_task = task
        coords = out_coords + dummies
        args = []
        for cmap, axes, (arg, ind), key in zip(
            coord_maps, concat_axes, argpairs, keys, strict=True
        ):
            if key not in task.dependencies:
                # FIXME: This feels like a bug
                continue
            if ind is None: