Repository: dask/dask
File: dask/blockwise.py (Python)

Code before:
    """

    def __init__(
        self,
        output,
        output_indices,
        dsk,
        indices,
        numblocks,
        concatenate=None,
        new_axes=None,
        io_subgraph=None,
    ):
        self.output = output
        self.output_indices = tuple(output_indices)
        self.io_subgraph = io_subgraph[1] if io_subgraph else None
        self.io_name = io_subgraph[0] if io_subgraph else None
        if not dsk:
            # If there is no `dsk` input, there must be an IO subgraph.
            if io_subgraph is None:
                raise ValueError("io_subgraph required if dsk is not supplied.")

            # Extract actual IO function for SubgraphCallable construction.
            # Wrap func in `PackedFunctionCall`, since it will receive
            # all arguments as a sigle (packed) tuple at run time.
            io_func = self.io_subgraph.get((self.io_name, 0), (None,))[0]
            ninds = 1 if isinstance(output_indices, str) else len(output_indices)
            dsk = {
                output: (
                    PackedFunctionCall(io_func),
                    *[blockwise_token(i) for i in range(ninds)],
                )
            }

        self.dsk = dsk
        self.indices = tuple(
            (name, tuple(ind) if ind is not None else ind) for name, ind in indices
        )
        self.numblocks = numblocks
        # optimize_blockwise won't merge where `concatenate` doesn't match, so
        # enforce a canonical value if there are no axes for reduction.
        if all(
            all(i in output_indices for i in ind)
            for name, ind in indices
            if ind is not None
        ):
            self.concatenate = None
        else:
            self.concatenate = concatenate
        self.new_axes = new_axes or {}


Review comment (suggestion):
We should document whatever assumptions we're making about the  structure of `io_subgraph`

Code after:
        within Blockwise.

    See Also
    --------
    dask.blockwise.blockwise
    dask.array.blockwise
    """

    def __init__(
        self,
        output,
        output_indices,
        dsk,
        indices,
        numblocks,
        concatenate=None,
        new_axes=None,
        io_subgraph=None,
    ):
        self.output = output
        self.output_indices = tuple(output_indices)
        self.io_subgraph = io_subgraph[1] if io_subgraph else None
        self.io_name = io_subgraph[0] if io_subgraph else None
        if not dsk:
            # If there is no `dsk` input, there must be an IO subgraph.
            if io_subgraph is None:
                raise ValueError("io_subgraph required if dsk is not supplied.")

            # Extract actual IO function for SubgraphCallable construction.
            # Wrap func in `PackedFunctionCall`, since it will receive
            # all arguments as a sigle (packed) tuple at run time.
            if self.io_subgraph:
                # We assume a 1-to-1 mapping between keys (i.e. tasks) and
                # chunks/partitions in `io_subgraph`, and assume the first
                # (callable) element is the same for all tasks.
                any_key = next(iter(self.io_subgraph))
                io_func = self.io_subgraph.get(any_key)[0]
            else:
                io_func = None
            ninds = 1 if isinstance(output_indices, str) else len(output_indices)
            dsk = {
                output: (
                    PackedFunctionCall(io_func),
                    *[blockwise_token(i) for i in range(ninds)],
                )
            }

        self.dsk = dsk
        self.indices = tuple(
            (name, tuple(ind) if ind is not None else ind) for name, ind in indices
        )