def topk(k, x):
    if x.ndim != 1:
        raise ValueError("Topk only works on arrays of one dimension")

    token = tokenize(k, x)
    name = 'chunk.topk-' + token
    dsk = dict(((name, i), (chunk.topk, k, key))
                for i, key in enumerate(x._keys()))
    name2 = 'topk-' + token
    dsk[(name2, 0)] = (getitem,
                        (np.sort, (np.concatenate, (list, list(dsk)))),
                        slice(-1, -k - 1, -1))
    chunks = ((k,),)

    return Array(merge(dsk, x.dask), name2, chunks, dtype=x.dtype)
The top k elements of an array
dataset/ETHPy150Open dask/dask/dask/array/core.py/topk