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

Code before:
        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]
        dtype = np.min_scalar_type(max(sorted_outblock_idx))
        sorted_outblock_idx = sorted_outblock_idx.astype(dtype)
        # 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),
                        interleave_none(full_slices, inblock),
                    ),
                    axis,

Review comment (suggestion):
We already know this:
```suggestion
        dtype = np.min_scalar_type(max_chunk_point_dimensions)
        sorted_outblock_idx = sorted_outblock_idx.astype(dtype, copy=False)
```

Code after:
            ]
        )

        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]
        dtype = np.min_scalar_type(max_chunk_point_dimensions)
        sorted_outblock_idx = sorted_outblock_idx.astype(dtype, copy=False)
        # 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),