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

Code before:
    def _subset_to_indexed_axes(iterable):
        for i, elem in enumerate(iterable):
            if i in axes:
                yield elem

    bounds2 = tuple(
        np.array(cached_cumsum(c, initial_zero=True))
        for c in _subset_to_indexed_axes(x.chunks)
    )
    axis = _get_axis(tuple(i if i in axes else None for i in range(x.ndim)))
    out_name = "vindex-merge-" + token

    # Now compute indices of each output element within each input block
    # The index is relative to the block, not the array.
    block_idxs = tuple(
        np.searchsorted(b, ind, side="right") - 1
        for b, ind in zip(bounds2, dict_indexes.values())
    )
    starts = (b[i] for i, b in zip(block_idxs, bounds2))
    inblock_idxs = np.broadcast_arrays(
        *(idx - start for idx, start in zip(dict_indexes.values(), starts))
    )
    for i, a in enumerate(inblock_idxs):
        if len(a) > 0:
            dtype = np.min_scalar_type(np.max(a, axis=None))
            inblock_idxs[i] = a.astype(dtype)

    chunks = [c for i, c in enumerate(x.chunks) if i not in axes]
    # determine number of points in one single output block.
    # Use the input chunk size to determine this.
    max_chunk_point_dimensions = reduce(
        mul, map(cached_max, _subset_to_indexed_axes(x.chunks))
    )

    n_chunks, remainder = divmod(npoints, max_chunk_point_dimensions)
    chunks.insert(
        0,
        (
            (max_chunk_point_dimensions,) * n_chunks
            + ((remainder,) if remainder > 0 else ())
            if npoints > 0
            else (0,)
        ),
    )
    chunks = tuple(chunks)

    if npoints > 0:
        other_blocks = product(
            *[
                range(len(c)) if i not in axes else [None]
                for i, c in enumerate(x.chunks)

Review comment (suggestion):
let's do this before broadcasting in the comprehension above and pass `copy=False` (FWIW)

Code after:
    def _subset_to_indexed_axes(iterable):
        for i, elem in enumerate(iterable):
            if i in axes:
                yield elem

    bounds2 = tuple(
        np.array(cached_cumsum(c, initial_zero=True))
        for c in _subset_to_indexed_axes(x.chunks)
    )
    axis = _get_axis(tuple(i if i in axes else None for i in range(x.ndim)))
    out_name = "vindex-merge-" + token

    # Now compute indices of each output element within each input block
    # The index is relative to the block, not the array.
    block_idxs = tuple(
        np.searchsorted(b, ind, side="right") - 1
        for b, ind in zip(bounds2, dict_indexes.values())
    )
    starts = (b[i] for i, b in zip(block_idxs, bounds2))
    inblock_idxs = []
    for idx, start in zip(dict_indexes.values(), starts):
        a = idx - start
        if len(a) > 0:
            dtype = np.min_scalar_type(np.max(a, axis=None))
            inblock_idxs.append(a.astype(dtype, copy=False))
        else:
            inblock_idxs.append(a)

    inblock_idxs = np.broadcast_arrays(*inblock_idxs)

    chunks = [c for i, c in enumerate(x.chunks) if i not in axes]
    # determine number of points in one single output block.
    # Use the input chunk size to determine this.
    max_chunk_point_dimensions = reduce(
        mul, map(cached_max, _subset_to_indexed_axes(x.chunks))
    )

    n_chunks, remainder = divmod(npoints, max_chunk_point_dimensions)
    chunks.insert(
        0,
        (
            (max_chunk_point_dimensions,) * n_chunks
            + ((remainder,) if remainder > 0 else ())
            if npoints > 0
            else (0,)
        ),
    )
    chunks = tuple(chunks)

    if npoints > 0:
        other_blocks = product(