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

Code before:
        # if input chunk (0,0) contributes to output chunks 0 and 2 then
        # (0,0) => (0, 0, 0) and (2, 0, 0).
        # The following is a groupby over output key by using ravel_multi_index to convert
        # the (output_block, *input_block) tuple to a unique code.
        ravel_shape = (
            n_chunks + 1,
            *(
                n
                for i, n in zip(flat_indexes, x.numblocks, strict=True)
                if i is not None
            ),
        )
        keys = np.ravel_multi_index([outblocks, *block_idxs], ravel_shape)
        unique_keys = np.unique(keys)

        name = "vindex-slice-" + token
        vindex_merge_name = "vindex-merge-" + token
        dsk = {}
        mask = np.empty(broadcast_shape, dtype=bool)
        for okey in other_blocks:
            merge_inputs = defaultdict(list)
            merge_indexer = defaultdict(list)
            for i, key in enumerate(unique_keys):
                outblock, *input_blocks = np.unravel_index(key, ravel_shape)
                np.equal(keys, key, out=mask)
                inblock = [_[mask] for _ in inblock_idxs]
                dsk[keyname(name, i, okey)] = (
                    _vindex_transpose,
                    (
                        _vindex_slice,
                        (x.name,) + interleave_none(okey, input_blocks),
                        interleave_none(full_slices, inblock),
                    ),
                    axis,
                )
                merge_inputs[outblock].append(keyname(name, i, okey))
                merge_indexer[outblock].append(outblock_idx[mask])

            for i in merge_inputs.keys():
                dsk[keyname(vindex_merge_name, i, okey)] = (
                    _vindex_merge,
                    merge_indexer[i],
                    merge_inputs[i],
                )

        result_1d = Array(
            HighLevelGraph.from_collections(out_name, dsk, dependencies=[x]),
            out_name,
            chunks,
            x.dtype,
            meta=x._meta,

Review comment (performance):
Ideally we'd figure out how to index before broadcasting and broadcast here but that will add a boatload of complexity

Code after:
                range(len(c)) if i not in axes else [None]
                for i, c in enumerate(x.chunks)
            ]
        )

        full_slices = [
            slice(None, None) if i not in axes else None for i in range(x.ndim)
        ]

        # The output is constructed as a new dimension and then reshaped
        # So the index of the output point is simply an `arange`
        outinds = np.arange(npoints).reshape(broadcast_shape)
        # Which output block is the point going to, and what is the index within that block?
        outblocks, outblock_idx = np.divmod(outinds, max_chunk_point_dimensions)

        # Now we try to be clever. We need to construct a graph where
        # if input chunk (0,0) contributes to output chunks 0 and 2 then
        # (0,0) => (0, 0, 0) and (2, 0, 0).
        # The following is a groupby over output key by using ravel_multi_index to convert
        # the (output_block, *input_block) tuple to a unique code.
        ravel_shape = (n_chunks + 1, *_subset_to_indexed_axes(x.numblocks))
        keys = np.ravel_multi_index([outblocks, *block_idxs], ravel_shape)
        # by sorting the data that needs to be inserted in the graph here,
        # we can slice in the hot loop below instead of using fancy indexing which will
        # always copy inside the hot loop.
        sortidx = np.argsort(keys, axis=None)
        sorted_keys = keys.flat[sortidx]  # flattens
        sorted_inblock_idxs = [_.flat[sortidx] for _ in inblock_idxs]
        sorted_outblock_idx = outblock_idx.flat[sortidx]
        # Determine the start and end of each unique key. We will loop over this
        flag = np.concatenate([[True], sorted_keys[1:] != sorted_keys[:-1], [True]])
        (key_bounds,) = flag.nonzero()

        name = "vindex-slice-" + token
        vindex_merge_name = "vindex-merge-" + token
        dsk = {}
        for okey in other_blocks:
            merge_inputs = defaultdict(list)
            merge_indexer = defaultdict(list)
            for i, (start, stop) in enumerate(
                zip(key_bounds[:-1], key_bounds[1:], strict=True)
            ):
                slicer = slice(start, stop)
                key = sorted_keys[start]
                outblock, *input_blocks = np.unravel_index(key, ravel_shape)
                inblock = [_[slicer] for _ in sorted_inblock_idxs]
                dsk[keyname(name, i, okey)] = (
                    _vindex_transpose,
                    (
                        _vindex_slice,
                        (x.name,) + interleave_none(okey, input_blocks),