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

Code before:
    Handles a single index by a single list

    Mimics ``np.take``

    >>> from pprint import pprint
    >>> chunks, dsk = take('y', 'x', [(20, 20, 20, 20)], [5, 1, 47, 3], axis=0)
    >>> chunks
    ((4,),)

    When list is sorted we still try to preserve properly sized chunks.

    >>> chunks, dsk = take('y', 'x', [(20, 20, 20, 20)], [1, 3, 5, 47], axis=0)
    >>> chunks
    ((4,),)

    """

    if not np.isnan(chunks[axis]).any():
        from dask.array._shuffle import _shuffle
        from dask.array.utils import arange_safe, asarray_safe

        arange = arange_safe(np.sum(chunks[axis]), like=index)
        if len(index) == len(arange) and np.abs(index - arange).sum() == 0:
            # TODO: This should be a real no-op, but the call stack is
            # too deep to do this efficiently for now
            chunk_tuples = list(product(*(range(len(c)) for i, c in enumerate(chunks))))
            graph = {
                (outname,) + c: Alias((outname,) + c, (inname,) + c)
                for c in chunk_tuples
            }
            return tuple(chunks), graph

        average_chunk_size = int(sum(chunks[axis]) / len(chunks[axis]))

        indexer = []
        index = asarray_safe(index, like=index)
        for i in range(0, len(index), average_chunk_size):
            indexer.append(index[i : i + average_chunk_size].tolist())

        token = (
            outname.split("-")[-1]
            if "-" in outname
            else tokenize(outname, chunks, index, axis)
        )
        chunks, graph = _shuffle(chunks, indexer, axis, inname, outname, token)
        return chunks, graph
    elif len(chunks[axis]) == 1:
        slices = [slice(None)] * len(chunks)
        slices[axis] = list(index)
        slices = tuple(slices)
        chunk_tuples = product(*(range(len(c)) for i, c in enumerate(chunks)))

Review comment (suggestion):
```suggestion
            chunk_tuples = product(*(range(len(c)) for i, c in enumerate(chunks)))
```

minor suggestion

Code after:
    Handles a single index by a single list

    Mimics ``np.take``

    >>> from pprint import pprint
    >>> chunks, dsk = take('y', 'x', [(20, 20, 20, 20)], [5, 1, 47, 3], axis=0)
    >>> chunks
    ((4,),)

    When list is sorted we still try to preserve properly sized chunks.

    >>> chunks, dsk = take('y', 'x', [(20, 20, 20, 20)], [1, 3, 5, 47], axis=0)
    >>> chunks
    ((4,),)

    """

    if not np.isnan(chunks[axis]).any():
        from dask.array._shuffle import _shuffle
        from dask.array.utils import arange_safe, asarray_safe

        arange = arange_safe(np.sum(chunks[axis]), like=index)
        if len(index) == len(arange) and np.abs(index - arange).sum() == 0:
            # TODO: This should be a real no-op, but the call stack is
            # too deep to do this efficiently for now
            chunk_tuples = product(*(range(len(c)) for i, c in enumerate(chunks)))
            graph = {
                (outname,) + c: Alias((outname,) + c, (inname,) + c)
                for c in chunk_tuples
            }
            return tuple(chunks), graph

        average_chunk_size = int(sum(chunks[axis]) / len(chunks[axis]))

        indexer = []
        index = asarray_safe(index, like=index)
        for i in range(0, len(index), average_chunk_size):
            indexer.append(index[i : i + average_chunk_size].tolist())

        token = (
            outname.split("-")[-1]
            if "-" in outname
            else tokenize(outname, chunks, index, axis)
        )
        chunks, graph = _shuffle(chunks, indexer, axis, inname, outname, token)
        return chunks, graph
    elif len(chunks[axis]) == 1:
        slices = [slice(None)] * len(chunks)
        slices[axis] = list(index)
        slices = tuple(slices)
        chunk_tuples = product(*(range(len(c)) for i, c in enumerate(chunks)))