dask: NEP-18: diag() fails working on CuPy array

Small sample to reproduce and traceback:

import cupy
import dask.array as da

x = cupy.random.random(5000)

d = da.from_array(x, chunks=(1000), asarray=False)

print(da.diag(d).compute())
Traceback (most recent call last):
  File "diag.py", line 8, in <module>
    print(da.diag(d).compute())
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/base.py", line 156, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/base.py", line 399, in compute
    return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/base.py", line 399, in <listcomp>
    return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/array/core.py", line 779, in finalize
    return concatenate3(results)
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/array/core.py", line 3499, in concatenate3
    return _concatenate2(arrays, axes=list(range(x.ndim)))
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/array/core.py", line 227, in _concatenate2
    arrays = [_concatenate2(a, axes=axes[1:]) for a in arrays]
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/array/core.py", line 227, in <listcomp>
    arrays = [_concatenate2(a, axes=axes[1:]) for a in arrays]
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/array/core.py", line 229, in _concatenate2
    return concatenate(arrays, axis=axes[0])
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/cupy/manipulation/join.py", line 49, in concatenate
    return core.concatenate_method(tup, axis)
  File "cupy/core/_routines_manipulation.pyx", line 560, in cupy.core._routines_manipulation.concatenate_method
  File "cupy/core/_routines_manipulation.pyx", line 573, in cupy.core._routines_manipulation.concatenate_method
TypeError: Only cupy arrays can be concatenated

A little debugging shows that Dask ends up creating NumPy temporary arrays, and later tries to concatenate CuPy with NumPy arrays.

@mrocklin

About this issue

  • Original URL
  • State: closed
  • Created 5 years ago
  • Comments: 21 (21 by maintainers)

Most upvoted comments

@TAdeJong this is being handled in https://github.com/dask/dask/pull/4543, in fact I got them to work last week, but that’s a long running PR, so it may take a little while to have integrated.

Indeed, I confirmed this on my side too. Thanks for pinging here @jakirkham !