Repository: dask/dask
Pull request: Modify mean chunk functions to return dicts rather than arrays
These functions return two array chunks each:

1.  The sum of the array
2.  The counts of the array

We need to return these as a single intermediate value.

Previously we did this by constructing an empty numpy array and then
assigning into it.  This doesn't work as well with numpy-like arrays
like cupy and sparse.

Now we return a dict containing two arrays.  This is nice in a few ways,
but does mean that we need to replicate the `concatenate=`
functionality within our chunk aggregation and combine functions, which
can be tricky.  This commit fails to do it correctly, but may be a good
start.

This is also somewhat suboptimal in a distributed setting because we no longer get special-cased handling of numpy arrays like we used to.  The scheduler will probably miscount the number of bytes that these objects hold and will probably not use ideal serialization and compression.  I *think* that this is probably ok because they're likely to be small, but it's worth pointing out.

After this we would like to implement the same fix for other functions,
notably the `moment_*` functions in this same reductions.py file.

Fixes https://github.com/dask/dask/issues/4481

Help or ideas with this would be welcome.  cc @pentschev @jakirkham 

- [ ] Tests added / passed
- [ ] Passes `flake8 dask`


Problem statement:
NEP-18: mean_chunk() object __array__ method not producing an array
Several Dask operations that utilize `mean_chunk()` from `dask/array/reductions.py` fail for Dask arrays created from non-NumPy (e.g., CuPy, sparse) arrays.

Some of the operations confirmed to fail are (including other non-core Dask projects):
* `dask.mean()`
* `dask.glm.algorithms.*`

A sample to reproduce the issue is given below, followed by its traceback.

```python
import cupy
import dask.array as da

x = cupy.random.random((5000, 1000))

d = da.from_array(x, chunks=(1000, 1000), asarray=False)

d.mean().compute()
```

```python
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  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 398, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/threaded.py", line 76, in get
    pack_exception=pack_exception, **kwargs)
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/local.py", line 459, in get_async
    raise_exception(exc, tb)
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/compatibility.py", line 112, in reraise
    raise exc
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/local.py", line 230, in execute_task
    result = _execute_task(task, data)
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/core.py", line 119, in _execute_task
    return func(*args2)
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/optimization.py", line 942, in __call__
    dict(zip(self.inkeys, args)))
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/core.py", line 149, in get
    result = _execute_task(task, cache)
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/core.py", line 119, in _execute_task
    return func(*args2)
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/compatibility.py", line 93, in apply
    return func(*args, **kwargs)
  File "/home/nfs/pentschev/.local/lib/python3.5/site-packages/dask/array/reductions.py", line 336, in mean_chunk
    result['n'] = n
ValueError: object __array__ method not producing an array
```

Hints from issue discussion:
Hrm, @jcrist do you have any suggestions here?  The challenge is that we're creating an `empty` numpy array and then assigning into it.  This makes things hard for non-Numpy implementations like sparse or cupy.

Two thoughts:

1.  We could see if dask array is returning a tuple rather than a record array.  My guess is that it will get upset
2.  We could use `empty_like` for our dispatch rather than `empty`.  This would allow the `__array_function__` protocol to work as desired.
We already have a dispatch for `empty`, found in `dask.array.reductions`, could cupy just register an implementation for that?

`empty_like` won't work, since you can't override the shape (unfortunately).
Right, following @jcrist 's suggestion this should be easy.  

@pentschev you probably want something like the following in `dask/array/backends.py`

```python
from .reductions import empty_lookup
empty_lookup.register(cupy.ndarray, np.cupy)
```
> ```python
> from .reductions import empty_lookup
> empty_lookup.register(cupy.ndarray, np.cupy)
> ```

Yes, that probably would solve the problem, but `cupy.empty()` doesn't support a list of tuples as `dtype`. Just as means of reproduction consider the following:

```python
In [25]: n = cupy.array([[1000000.]])

In [26]: n
Out[26]: array([[1000000.]])

In [27]: n.shape
Out[27]: (1, 1)

In [28]: cupy.empty(n.shape, dtype=[('total', total.dtype), ('n', total.dtype)])
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-28-ea9bfab84d63> in <module>
----> 1 cupy.empty(n.shape, dtype=[('total', total.dtype), ('n', total.dtype)])

~/.local/lib/python3.5/site-packages/cupy/creation/basic.py in empty(shape,
dtype, order)
     20 
     21     """
---> 22     return cupy.ndarray(shape, dtype, order=order)
     23 
     24 

~/src/cupy/cupy/core/core.pyx in cupy.core.core.ndarray.__init__()

~/src/cupy/cupy/core/_dtype.pyx in cupy.core._dtype.get_dtype_with_itemsize()

TypeError: unhashable type: 'list'
```

I'm not sure whether this is a bug or a missing feature in CuPy. But as far as I can tell, this empty array is only used for internal Dask computing, so an alternative would be to replace it by something else, since it's only used to hold two values by their names, just like a hash table.
Lets raise the list-of-tuple dtype issue on the cupy issue tracker to see
if they can add it.

In the meantime we might look to see if we can avoid using record arrays
(these are what list-of-tuple dtype arrays are called).  I can take a quick
look at this today.

On Thu, Feb 14, 2019 at 1:54 AM Peter Andreas Entschev <
notifications@github.com> wrote:

> from .reductions import empty_lookup
> empty_lookup.register(cupy.ndarray, np.cupy)
>
> Yes, that probably would solve the problem, but cupy.empty() doesn't
> support a list of tuples as dtype. Just as means of reproduction consider
> the following:
>
> In [25]: n = cupy.array([[1000000.]])
>
> In [26]: n
> Out[26]: array([[1000000.]])
>
> In [27]: n.shape
> Out[27]: (1, 1)
>
> In [28]: cupy.empty(n.shape, dtype=[('total', total.dtype), ('n', total.dtype)])---------------------------------------------------------------------------TypeError                                 Traceback (most recent call last)<ipython-input-28-ea9bfab84d63> in <module>----> 1 cupy.empty(n.shape, dtype=[('total', total.dtype), ('n', total.dtype)])
> ~/.local/lib/python3.5/site-packages/cupy/creation/basic.py in empty(shape,
> dtype, order)
>      20
>      21     """---> 22     return cupy.ndarray(shape, dtype, order=order)     23      24 ~/src/cupy/cupy/core/core.pyx in cupy.core.core.ndarray.__init__()~/src/cupy/cupy/core/_dtype.pyx in cupy.core._dtype.get_dtype_with_itemsize()TypeError: unhashable type: 'list'
>
> I'm not sure whether this is a bug or a missing feature in CuPy. But as
> far as I can tell, this empty array is only used for internal Dask
> computing, so an alternative would be to replace it by something else,
> since it's only used to hold two values by their names, just like a hash
> table.
>
> —
> You are receiving this because you commented.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/4481#issuecomment-463564701>, or mute
> the thread
> <https://github.com/notifications/unsubscribe-auth/AASszLBDBppeOJ1b2xUH9_1LCGsLWw1Gks5vNTJagaJpZM4a5eUj>
> .
>

@mrocklin I'm opening the issue on CuPy's tracker. If you find an appropriate way of avoiding the use of record arrays, I could also work on the implementation if you'd like, just let me know.
Raised here: https://github.com/cupy/cupy/issues/2031

It looks like we can resolve this on our end just by adding a `np.dtype(...)` call around the list.
So the following change is probably needed and welcome on the dask side

```diff
-   result = empty(n.shape, dtype=[('total', total.dtype), ('n', n.dtype)])
+   result = empty(n.shape, dtype=np.dtype([('total', total.dtype), ('n', n.dtype)]))
```
Yes, that should work. I can try it out now and submit a PR, or are you working on this already @mrocklin ?
All yours

Commit message: Modify mean chunk functions to return dicts rather than arrays

These functions return two array chunks each:

1.  The sum of the array
2.  The counts of the array

We need to return these as a single intermediate value.

Previously we did this by constructing an empty numpy array and then
assigning into it.  This doesn't work as well with numpy-like arrays
like cupy and sparse.

Now we return a dict containing two arrays.  This is nice in a few ways,
but does mean that we need to replicate the `concatenate=`
functionality within our chunk aggregation and combine functions, which
can be tricky.  This commit fails to do it correctly, but may be a good
start.

After this we would like to implement the same fix for other functions,
notably the `moment_*` functions in this same reductions.py file.

Fixes https://github.com/dask/dask/issues/4481

Patch to review:
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
index 947e24d8819..0b0c14dc98b 100644
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -22,7 +22,7 @@
 from ..compatibility import getargspec, builtins
 from ..base import tokenize
 from ..highlevelgraph import HighLevelGraph
-from ..utils import ignoring, funcname, Dispatch
+from ..utils import ignoring, funcname, Dispatch, deepmap
 from .. import config
 
 # Generic functions to support chunks of different types
@@ -331,26 +331,27 @@ def nannumel(x, **kwargs):
 def mean_chunk(x, sum=chunk.sum, numel=numel, dtype='f8', **kwargs):
     n = numel(x, dtype=dtype, **kwargs)
     total = sum(x, dtype=dtype, **kwargs)
-    empty = empty_lookup.dispatch(type(n))
-    result = empty(n.shape, dtype=[('total', total.dtype), ('n', n.dtype)])
-    result['n'] = n
-    result['total'] = total
-    return result
+    return {'n': n, 'total': total}
 
 
-def mean_combine(pair, sum=chunk.sum, numel=numel, dtype='f8', **kwargs):
-    n = sum(pair['n'], **kwargs)
-    total = sum(pair['total'], **kwargs)
-    empty = empty_lookup.dispatch(type(n))
-    result = empty(n.shape, dtype=pair.dtype)
-    result['n'] = n
-    result['total'] = total
-    return result
+def mean_combine(pairs, sum=chunk.sum, numel=numel, dtype='f8', axis=None, **kwargs):
+    if not isinstance(pairs, list):
+        pairs = [pairs]
+    ns = deepmap(lambda pair: pair['n'], pairs)
+    totals = deepmap(lambda pair: pair['total'], pairs)
+    n = _concatenate2(ns, axes=axis).sum(axis=axis, **kwargs)
+    total = _concatenate2(totals, axes=axis).sum(axis=axis, **kwargs)
+    return {'n': n, 'total': total}
 
 
-def mean_agg(pair, dtype='f8', **kwargs):
-    return divide(pair['total'].sum(dtype=dtype, **kwargs),
-                  pair['n'].sum(dtype=dtype, **kwargs), dtype=dtype)
+def mean_agg(pairs, dtype='f8', axis=None, **kwargs):
+    ns = deepmap(lambda pair: pair['n'], pairs)
+    totals = deepmap(lambda pair: pair['total'], pairs)
+    n = _concatenate2(ns, axes=axis).sum(axis=axis, **kwargs)
+    total = _concatenate2(totals, axes=axis).sum(axis=axis, **kwargs)
+
+    return divide(total.sum(dtype=dtype, **kwargs),
+                  n.sum(dtype=dtype, **kwargs), dtype=dtype)
 
 
 @wraps(chunk.mean)
@@ -361,7 +362,7 @@ def mean(a, axis=None, dtype=None, keepdims=False, split_every=None, out=None):
         dt = getattr(np.mean(np.empty(shape=(1,), dtype=a.dtype)), 'dtype', object)
     return reduction(a, mean_chunk, mean_agg, axis=axis, keepdims=keepdims,
                      dtype=dt, split_every=split_every, combine=mean_combine,
-                     out=out)
+                     out=out, concatenate=False)
 
 
 def nanmean(a, axis=None, dtype=None, keepdims=False, split_every=None,
@@ -373,6 +374,7 @@ def nanmean(a, axis=None, dtype=None, keepdims=False, split_every=None,
     return reduction(a, partial(mean_chunk, sum=chunk.nansum, numel=nannumel),
                      mean_agg, axis=axis, keepdims=keepdims, dtype=dt,
                      split_every=split_every, out=out,
+                     concatenate=False,
                      combine=partial(mean_combine, sum=chunk.nansum, numel=nannumel))
 
 


Code review comments:
diff_hunk: @@ -331,26 +331,27 @@ def nannumel(x, **kwargs):
 def mean_chunk(x, sum=chunk.sum, numel=numel, dtype='f8', **kwargs):
     n = numel(x, dtype=dtype, **kwargs)
     total = sum(x, dtype=dtype, **kwargs)
-    empty = empty_lookup.dispatch(type(n))
-    result = empty(n.shape, dtype=[('total', total.dtype), ('n', n.dtype)])
-    result['n'] = n
-    result['total'] = total
-    return result
+    return {'n': n, 'total': total}
 
 
-def mean_combine(pair, sum=chunk.sum, numel=numel, dtype='f8', **kwargs):
-    n = sum(pair['n'], **kwargs)
-    total = sum(pair['total'], **kwargs)
-    empty = empty_lookup.dispatch(type(n))
-    result = empty(n.shape, dtype=pair.dtype)
-    result['n'] = n
-    result['total'] = total
-    return result
+def mean_combine(pairs, sum=chunk.sum, numel=numel, dtype='f8', axis=None, **kwargs):
+    if not isinstance(pairs, list):
+        pairs = [pairs]
+    ns = deepmap(lambda pair: pair['n'], pairs)
+    totals = deepmap(lambda pair: pair['total'], pairs)
+    n = _concatenate2(ns, axes=axis).sum(axis=axis, **kwargs)
+    total = _concatenate2(totals, axes=axis).sum(axis=axis, **kwargs)
+    return {'n': n, 'total': total}
 
 
-def mean_agg(pair, dtype='f8', **kwargs):
-    return divide(pair['total'].sum(dtype=dtype, **kwargs),
-                  pair['n'].sum(dtype=dtype, **kwargs), dtype=dtype)
+def mean_agg(pairs, dtype='f8', axis=None, **kwargs):
+    ns = deepmap(lambda pair: pair['n'], pairs)
+    totals = deepmap(lambda pair: pair['total'], pairs)
+    n = _concatenate2(ns, axes=axis).sum(axis=axis, **kwargs)
+    total = _concatenate2(totals, axes=axis).sum(axis=axis, **kwargs)
+
+    return divide(total.sum(dtype=dtype, **kwargs),
+                  n.sum(dtype=dtype, **kwargs), dtype=dtype)
path: dask/array/reductions.py
text: @user1:
Found the mistake, you're calling `.sum()` twice for both `total` and `n`. Get rid of one and you should be good to go!