Using multiple backends

[pentschev]

Following the discussion we had in https://github.com/numpy/numpy/issues/13831, I was trying uarray out, and wanted to test multiple backends.

@peterbell10 commented that it should be with ua.set_backend(inner_array_backend), ua.set_backend(outer_array_backend):, and while I couldn't find any examples, this test suggests it should be the other way around.

Nevertheless, I'm trying to create Dask arrays that use CuPy as a backend, and I tried both orders of inner/outer arrays. It would also be interesting to clarify somewhere what's the inner and outer arrays, I assume in the test I'm trying, CuPy would be inner and Dask would be outer, but feel free to correct me if I'm looking it in the wrong way. Here's what I could get:

>>> import uarray as ua
>>> import unumpy as np
>>> import unumpy.dask_backend as dask_backend
>>> import unumpy.cupy_backend as cupy_backend
>>>
>>> with ua.set_backend(cupy_backend), ua.set_backend(dask_backend):
...     a = np.ones((2, 2))
...     print(a)
...     print(type(a), type(a._meta))
...     print(a.compute())
...     print(type(a.compute()))
...
dask.array<ones, shape=(2, 2), dtype=float64, chunksize=(2, 2)>
<class 'dask.array.core.Array'> <class 'numpy.ndarray'>
[[1. 1.]
 [1. 1.]]
<class 'numpy.ndarray'>

I expected in this case to have type(a._meta) and type(a.compute()) to be CuPy arrays, but got NumPy. Am I doing something wrong?

If I invert the cupy_backend/dask_backend orders, then a is a CuPy array, so it can't wrap other arrays.

cc @hameerabbasi

[peterbell10]

@peterbell10 commented that it should be with ua.set_backend(inner_array_backend), ua.set_backend(outer_array_backend):, and while I couldn't find any examples, this test suggests it should be the other way around.

I think it's just that the variable names in that test are backwards: be_inner actually delegates to be_outer which has the concrete "implementation" (returning obj).

I expected in this case to have type(a._meta) and type(a.compute()) to be CuPy arrays, but got NumPy. Am I doing something wrong?

I think to properly support meta-arrays, the backend need be written with uarray support in mind. At the moment unumpy.ones will just call da.ones which isn't uarray-aware so never calls into the CuPy backend.

[pentschev]

@peterbell10 thanks for the quick answer.

I think it's just that the variable names in that test are backwards: be_inner actually delegates to be_outer which has the concrete "implementation" (returning obj).

I see, thanks for clarifying.

I think to properly support meta-arrays, the backend need be written with uarray support in mind. At the moment unumpy.ones will just call da.ones which isn't uarray-aware so never calls into the CuPy backend.

So once Dask is uarray-aware, the example above would work as I thought it would, is that correct?

[peterbell10]

So once Dask is uarray-aware, the example above would work as I thought it would, is that correct?

Yes, I think you're correct.

[hameerabbasi]

@pentschev To elaborate, you would have to do basically what was done in the blogpost, and replace numpy with unumpy everywhere in the Dask codebase.

This will require an effort to define Multimethods for unumpy. Quite a few already exist.

[pentschev]

@pentschev To elaborate, you would have to do basically what was done in the blogpost, and replace numpy with unumpy everywhere in the Dask codebase.

What blogpost are you referring to?

This will require an effort to define Multimethods for unumpy. Quite a few already exist.

Would you mind pointing me to some of them?

[hameerabbasi]

Would you mind pointing me to some of them?

https://github.com/Quansight-Labs/uarray/blob/master/unumpy/multimethods.py contain all the multimethods.

What blogpost are you referring to?

https://labs.quansight.org/blog/2019/07/uarray-update-api-changes-overhead-and-comparison-to-__array_function__/

I updated it with a composition example using XArray and Dask that illustrates the principle.

[pentschev]

Great, thanks @hameerabbasi, I will try to take a look at this soon!

[hameerabbasi]

Closing this as the original question has been answered.