Repository: dask/dask

Issue:
[BUG] dask.array.einsum doesn't support CuPy backend
Running into numpy array related issues when trying to use cupy arrays with einsum. Looks like it's perhaps due to still getting to the `asarray` branching logic which appears to be set to True somewhere internally, despite the cupy inputs. Perhaps it's being reset or not fully passed down from the original arrays to an intermediate array.

```python
import cupy as cp
import dask.array as da
​
A = cp.array([0., 1, 2])
​
B = cp.array([[ 0,  1.,  2,  3],
              [ 4,  5,  6,  7],
              [ 8,  9, 10, 11]])
​
dA = da.from_array(A, asarray=False)
dB = da.from_array(B, asarray=False)
​
res = da.einsum('i,ij->i', dA, dB).compute()
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-147-06b849502129> in <module>
     11 dB = da.from_array(B, asarray=False)
     12 
---> 13 res = da.einsum('i,ij->i', dA, dB).compute()

/conda/envs/rapids/lib/python3.7/site-packages/dask/base.py in compute(self, **kwargs)
    154         dask.base.compute
    155         """
--> 156         (result,) = compute(self, traverse=False, **kwargs)
    157         return result
    158 

/conda/envs/rapids/lib/python3.7/site-packages/dask/base.py in compute(*args, **kwargs)
    396     keys = [x.__dask_keys__() for x in collections]
    397     postcomputes = [x.__dask_postcompute__() for x in collections]
--> 398     results = schedule(dsk, keys, **kwargs)
    399     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    400 

/conda/envs/rapids/lib/python3.7/site-packages/dask/threaded.py in get(dsk, result, cache, num_workers, pool, **kwargs)
     74     results = get_async(pool.apply_async, len(pool._pool), dsk, result,
     75                         cache=cache, get_id=_thread_get_id,
---> 76                         pack_exception=pack_exception, **kwargs)
     77 
     78     # Cleanup pools associated to dead threads

/conda/envs/rapids/lib/python3.7/site-packages/dask/local.py in get_async(apply_async, num_workers, dsk, result, cache, get_id, rerun_exceptions_locally, pack_exception, raise_exception, callbacks, dumps, loads, **kwargs)
    460                         _execute_task(task, data)  # Re-execute locally
    461                     else:
--> 462                         raise_exception(exc, tb)
    463                 res, worker_id = loads(res_info)
    464                 state['cache'][key] = res

/conda/envs/rapids/lib/python3.7/site-packages/dask/compatibility.py in reraise(exc, tb)
    110         if exc.__traceback__ is not tb:
    111             raise exc.with_traceback(tb)
--> 112         raise exc
    113 
    114     import pickle as cPickle

/conda/envs/rapids/lib/python3.7/site-packages/dask/local.py in execute_task(key, task_info, dumps, loads, get_id, pack_exception)
    228     try:
    229         task, data = loads(task_info)
--> 230         result = _execute_task(task, data)
    231         id = get_id()
    232         result = dumps((result, id))

/conda/envs/rapids/lib/python3.7/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    116     elif istask(arg):
    117         func, args = arg[0], arg[1:]
--> 118         args2 = [_execute_task(a, cache) for a in args]
    119         return func(*args2)
    120     elif not ishashable(arg):

/conda/envs/rapids/lib/python3.7/site-packages/dask/core.py in <listcomp>(.0)
    116     elif istask(arg):
    117         func, args = arg[0], arg[1:]
--> 118         args2 = [_execute_task(a, cache) for a in args]
    119         return func(*args2)
    120     elif not ishashable(arg):

/conda/envs/rapids/lib/python3.7/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    117         func, args = arg[0], arg[1:]
    118         args2 = [_execute_task(a, cache) for a in args]
--> 119         return func(*args2)
    120     elif not ishashable(arg):
    121         return arg

/conda/envs/rapids/lib/python3.7/site-packages/dask/optimization.py in __call__(self, *args)
    940                              % (len(self.inkeys), len(args)))
    941         return core.get(self.dsk, self.outkey,
--> 942                         dict(zip(self.inkeys, args)))
    943 
    944     def __reduce__(self):

/conda/envs/rapids/lib/python3.7/site-packages/dask/core.py in get(dsk, out, cache)
    147     for key in toposort(dsk):
    148         task = dsk[key]
--> 149         result = _execute_task(task, cache)
    150         cache[key] = result
    151     result = _execute_task(out, cache)

/conda/envs/rapids/lib/python3.7/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    117         func, args = arg[0], arg[1:]
    118         args2 = [_execute_task(a, cache) for a in args]
--> 119         return func(*args2)
    120     elif not ishashable(arg):
    121         return arg

/conda/envs/rapids/lib/python3.7/site-packages/dask/compatibility.py in apply(func, args, kwargs)
     91     def apply(func, args, kwargs=None):
     92         if kwargs:
---> 93             return func(*args, **kwargs)
     94         else:
     95             return func(*args)

/conda/envs/rapids/lib/python3.7/site-packages/dask/array/einsumfuncs.py in chunk_einsum(*operands, **kwargs)
     17     dtype = kwargs.pop('kernel_dtype')
     18     einsum = einsum_lookup.dispatch(type(operands[0]))
---> 19     chunk = einsum(subscripts, *operands, dtype=dtype, **kwargs)
     20 
     21     # Avoid concatenate=True in blockwise by adding 1's

/conda/envs/rapids/lib/python3.7/site-packages/numpy/core/einsumfunc.py in einsum(*operands, **kwargs)
   1344     # If no optimization, run pure einsum
   1345     if optimize_arg is False:
-> 1346         return c_einsum(*operands, **kwargs)
   1347 
   1348     valid_einsum_kwargs = ['out', 'dtype', 'order', 'casting']

ValueError: object __array__ method not producing an array
```

Some initial debugging
```python
> /conda/envs/rapids/lib/python3.7/site-packages/numpy/core/einsumfunc.py(1346)einsum()
   1344     # If no optimization, run pure einsum
   1345     if optimize_arg is False:
-> 1346         return c_einsum(*operands, **kwargs)
   1347 
   1348     valid_einsum_kwargs = ['out', 'dtype', 'order', 'casting']

ipdb>  type(operands[1])
<class 'cupy.core.core.ndarray'>
ipdb>  type(operands[2])
<class 'cupy.core.core.ndarray'>
```


Discussion:
@beckernick it seems that you don't have the `NUMPY_EXPERIMENTAL_ARRAY_FUNCTION=1` environment variable set, which is required with NumPy 1.16.x. Regardless, it still doesn't work when it's enabled, I've added this issue to our tracked in #4731. How urgently do you need this?
Thanks @pentschev . I don't think this is incredibly urgent (from my limited perspective), as the actual einstein summation usage by the user can probably be rewritten without too much effort with the appropriate indexing/transposing.

Will defer to @jakirkham in terms of general importance of this functionality versus other priorities (outside of this specific circumstance in which I noticed this).
@pentschev, do you have a sense of how difficult this might be to fix?
I was checking it now, and the only thing we need to add CuPy support for the newly introduced `_implementation` NumPy attribute https://github.com/numpy/numpy/pull/13627, that allows functions such as CuPy to fallback to the NumPy implementation when they do not wish to implement it themselves, which was not possible with `__array_function__` before. I will submit a PR to CuPy to make use of that feature.
Ah ok. That should be straightforward then. Also should help with a few use cases (including this one). Thanks for looking into it Peter. 🙂

Resolving patch:
diff --git a/dask/array/einsumfuncs.py b/dask/array/einsumfuncs.py
--- a/dask/array/einsumfuncs.py
+++ b/dask/array/einsumfuncs.py
@@ -194,14 +194,9 @@ def parse_einsum_input(operands):
 
 @derived_from(np)
 def einsum(*operands, **kwargs):
-    casting = kwargs.pop('casting', 'safe')
     dtype = kwargs.pop('dtype', None)
     optimize = kwargs.pop('optimize', False)
-    order = kwargs.pop('order', 'K')
     split_every = kwargs.pop('split_every', None)
-    if kwargs:
-        raise TypeError("einsum() got unexpected keyword "
-                        "argument(s) %s" % ",".join(kwargs))
 
     einsum_dtype = dtype
 
@@ -237,8 +232,8 @@ def einsum(*operands, **kwargs):
                        adjust_chunks={ind: 1 for ind in contract_inds}, dtype=dtype,
                        # np.einsum parameters
                        subscripts=subscripts, kernel_dtype=einsum_dtype,
-                       ncontract_inds=ncontract_inds, order=order,
-                       casting=casting, optimize=optimize)
+                       ncontract_inds=ncontract_inds,
+                       optimize=optimize, **kwargs)
 
     # Now reduce over any extra contraction dimensions
     if ncontract_inds > 0:


Test patch:
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -1529,7 +1529,7 @@ def test_einsum_split_every(split_every):
 def test_einsum_invalid_args():
     _, da_inputs = _numpy_and_dask_inputs('a')
     with pytest.raises(TypeError):
-        da.einsum('a', *da_inputs, foo=1, bar=2)
+        da.einsum('a', *da_inputs, foo=1, bar=2).compute()
 
 
 def test_einsum_broadcasting_contraction():