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

Code before:
def slicing_plan(chunks, index):
    """Construct a plan to slice chunks with the given index

    Parameters
    ----------
    chunks : Tuple[int]
        One dimensions worth of chunking information
    index : np.ndarray[int]
        The index passed to slice on that dimension

    Returns
    -------
    out : List[Tuple[int, np.ndarray]]
        A list of chunk/sub-index pairs corresponding to each output chunk
    """
    from dask.array.utils import asarray_safe

    if not is_arraylike(index):
        index = np.asanyarray(index)
    cum_chunks = cached_cumsum(chunks)

    cum_chunks = asarray_safe(cum_chunks, like=index)
    # this dispactches to the array library
    chunk_locations = np.searchsorted(cum_chunks, index, side="right")

    # but we need chunk_locations as python ints for getitem calls downstream
    chunk_locations = chunk_locations.tolist()
    where = np.where(np.diff(chunk_locations))[0] + 1

    extra = asarray_safe([0], like=where)
    c_loc = asarray_safe([len(chunk_locations)], like=where)
    where = np.concatenate([extra, where, c_loc])

    out = []
    for i in range(len(where) - 1):
        sub_index = index[where[i] : where[i + 1]]
        chunk = chunk_locations[where[i]]
        if chunk > 0:
            sub_index = sub_index - cum_chunks[chunk - 1]
        out.append((chunk, sub_index))

    return out


def take(outname, inname, chunks, index, itemsize, axis=0):
    """Index array with an iterable of index

    Handles a single index by a single list

    Mimics ``np.take``


Review comment (suggestion):
cum_chunks was int64 before (in numpy 1 too); this PR changes it to uint64 (or whatever the end user passed as a slicer).
This fixes the dtype arithmetics on line 586 below, which used to subtract uint64 - int64, which in numpy 1 returns uint64 but in numpy 2 returns float.

Code after:
def slicing_plan(chunks, index):
    """Construct a plan to slice chunks with the given index

    Parameters
    ----------
    chunks : Tuple[int]
        One dimensions worth of chunking information
    index : np.ndarray[int]
        The index passed to slice on that dimension

    Returns
    -------
    out : List[Tuple[int, np.ndarray]]
        A list of chunk/sub-index pairs corresponding to each output chunk
    """
    from dask.array.utils import asarray_safe

    if not is_arraylike(index):
        index = np.asanyarray(index)

    cum_chunks_tup = cached_cumsum(chunks)
    cum_chunks = asarray_safe(cum_chunks_tup, like=index)
    if cum_chunks.dtype.kind != "f":  # Don't cast NaN chunks to int
        # This is important when index.dtype=uint64 (or uint32 on 32-bit hosts) to
        # prevent accidental automatic casting during `index - cum_chunks` below
        cum_chunks = cum_chunks.astype(index.dtype)

    # this dispactches to the array library
    chunk_locations = np.searchsorted(cum_chunks, index, side="right")

    # but we need chunk_locations as python ints for getitem calls downstream
    chunk_locations = chunk_locations.tolist()
    where = np.where(np.diff(chunk_locations))[0] + 1

    extra = asarray_safe([0], like=where)
    c_loc = asarray_safe([len(chunk_locations)], like=where)
    where = np.concatenate([extra, where, c_loc])

    out = []
    for i in range(len(where) - 1):
        sub_index = index[where[i] : where[i + 1]]
        chunk = chunk_locations[where[i]]
        if chunk > 0:
            sub_index = sub_index - cum_chunks[chunk - 1]
        out.append((chunk, sub_index))

    return out


def take(outname, inname, chunks, index, itemsize, axis=0):
    """Index array with an iterable of index