Interop with Dask #22
Description
I am opening this issue to track how xnd/gumath could work with Dask based on talking with @mrocklin.
We can use the dask.array.from_array on an xnd object, to have Dask chunk up that array and execute operations in parallel. The requirements dask has for an array like object are listed here.
We were going back and forth on whether the numpy-ish API should be implemented on the xnd object directly or if we should create a wrapper class, like this, that adds numpy methods.
The most basic requirement is to have shape and dtype attributes on the object. We can forward those from the type attribute:
In [1]: from xnd import xnd
In [3]: x = xnd(list(range(10)))
In [4]: import dask.array as da
In [7]: x.shape = x.type.shape
In [15]: x.dtype = str(x.type.hidden_dtype)
In [17]: d = da.from_array(x, chunks=(5,))
In [19]: d.sum()
Out[19]: dask.array<sum-aggregate, shape=(), dtype=int64, chunksize=()>
In [21]: d.sum().compute()
Out[21]: 45
In [23]: np.exp(d).sum().compute()
Out[23]: 12818.308050524603
In [24]: d[::3]
Out[24]: dask.array<getitem, shape=(4,), dtype=int64, chunksize=(2,)>
In [25]: d[::3].compute()
Out[25]: array([0, 3, 6, 9])We could also implement __array_ufunc__ so that when we call np.sin on an xnd array, we could get back an xnd array by using the gumath.sin function. This is how cupy implements it.
We can also create a concat and register it with dask.array.core.concatenate_lookup. I don't think a concat gufunc exists yet.