array-api-compat
array-api-compat copied to clipboard
[WIP] ENH: dask+cupy, dask+sparse etc. namespaces
Very early prototype inspired by https://github.com/mdhaber/marray/
FYI @mdhaber @lucascolley
Demo
>>> import dask.array as da, cupy as cp
>>> from array_api_compat import array_namespace
>>> x = da.from_array(cp.asarray(0))
>>> xp = array_namespace(x)
>>> xp
<module 'array_api_compat.dask.array.cupy'>
>>> y = xp.arange(10)
>>> y
dask.array<arange, shape=(10,), dtype=int64, chunksize=(10,), chunktype=cupy.ndarray>
>>> type(y.compute())
cupy.ndarray
>>> array_namespace(y) is xp
True