Dask indexing problem with cupy

Hi, I want to index an array and then set those values to something.

So, this is the example that DOES NOT work.

import dask.array as da
import dask

dask.config.set({"array.backend": "cupy"})

data = da.random.random((20000,3))
condition = data[:,0] > 0.5
idx = da.flatnonzero(condition).compute()
vis[idx] *= -1

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/miguel/.conda/envs/pyralysis-env/lib/python3.12/site-packages/dask/array/core.py", line 1926, in __setitem__
    dsk = setitem_array(out, self, key, value)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/miguel/.conda/envs/pyralysis-env/lib/python3.12/site-packages/dask/array/slicing.py", line 1843, in setitem_array
    raise ValueError(
ValueError: could not broadcast input array from shape(9961, 3) into shape (3,)

However, if I try this with numpy as backend, it does work. It also works using numpy and cupywithout dask.

Any recommendation for this?

Hi @miguelcarcamov,

I can confirm I can reproduce your issue with Dask and Cupy, while it works using Numpy.

I think you should open an issue about this in Dask github issue tracker!