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

Code before:
                    if k[0] == self.output:
                        key_deps[k] = set()

            self._cached_dict = {
                "dsk": dsk,
                "basic_layer": BasicLayer(dsk, key_deps, non_blockwise_keys),
            }
        return self._cached_dict["dsk"]

    def get_output_keys(self):

        # TODO: Handle N-D Collections and more-complex
        # tensor operations.

        # Only deal with 1-D collections (for now).
        # Otherwise, we allow dict materialization.
        if len(self.output_indices) != 1:
            return super().get_output_keys()

        # Check inputs.
        # Only deal with 1-to-1 input-output collection
        # mapping (for now).
        input_cnt = 0
        in_name = None
        for _name, _ind in self.indices:
            if _ind is not None:
                if len(_ind) != 1 or input_cnt:
                    return super().get_output_keys()
                in_name = _name
                input_cnt += 1

        # At this point, we can assume:
        # - Input and output indices are aligned
        # - Collection is 1-D
        # - Only one input collection
        return {(self.output, p) for p in range(self.numblocks[in_name][0])}

    def __getitem__(self, key):
        return self._dict[key]

    def __iter__(self):
        return iter(self._dict)

    def __len__(self):
        return int(np.prod(list(self._out_numblocks().values())))

    def _out_numblocks(self):
        d = {}
        out_d = {}
        indices = {k: v for k, v in self.indices if v is not None}
        for k, v in self.numblocks.items():

Review comment (suggestion):
Could you open up an issue with checkboxes to track follow-up work we'll want to do to ensure we're using using IO-subgraphs throughout the codebase

Code after:
                # This is an IO layer.
                for k in dsk:
                    io_key = (self.io_name,) + tuple([k[i] for i in range(1, len(k))])
                    if io_key in dsk[k]:
                        # Inject IO-function arguments into the blockwise graph
                        # as a single (packed) tuple.
                        io_item = self.io_subgraph.get(io_key)
                        io_item = list(io_item[1:]) if len(io_item) > 1 else []
                        new_task = [io_item if v == io_key else v for v in dsk[k]]
                        dsk[k] = tuple(new_task)

                # Clear IO "placeholder" dependencies
                for k in key_deps:
                    if k[0] == self.output:
                        key_deps[k] = set()

            self._cached_dict = {
                "dsk": dsk,
                "basic_layer": BasicLayer(dsk, key_deps, non_blockwise_keys),
            }
        return self._cached_dict["dsk"]

    def get_output_keys(self):

        # TODO: Handle N-D Collections and more-complex
        # tensor operations.

        # Only deal with 1-D collections (for now).
        # Otherwise, we allow dict materialization.
        if len(self.output_indices) != 1:
            return super().get_output_keys()

        # Check inputs.
        # Only deal with 1-to-1 input-output collection
        # mapping (for now).
        input_cnt = 0
        in_name = None
        for _name, _ind in self.indices:
            if _ind is not None:
                if len(_ind) != 1 or input_cnt:
                    return super().get_output_keys()
                in_name = _name
                input_cnt += 1

        # At this point, we can assume:
        # - Input and output indices are aligned
        # - Collection is 1-D
        # - Only one input collection
        return {(self.output, p) for p in range(self.numblocks[in_name][0])}

    def __getitem__(self, key):