Repository: dask/dask
Pull request: Fixed mean() and moment() on sparse arrays
For proper creation of sparse auxiliary arrays, a dispatcher for ones_like has
been added. However, sparse arrays created from ones_like cause sparse
reductions to fail, as it would result in a dense array. To fix that, the arrays
have first to be converted to dense before, which ends up requiring a new
dispatcher for that purpose, and for already dense arrays must be bypassed, that
is here implemented as numpy.ndarray.view() to simply return a shallow copy of
the array.

This commit fixes https://github.com/dask/dask/issues/4523 and adds the tests
suggested in that issue.

@mrocklin 

- [ ] Tests added / passed
- [ ] Passes `flake8 dask`

Problem statement:
Some functions fail on Dask arrays from sparse
Short sample for reproduction and traceback follow.

```python
import numpy as np
import sparse
import dask.array as da
from numpy.testing import assert_array_equal

x = da.random.random((2, 3, 4), chunks=(1, 2, 2))
x[x < 0.8] = 0

y = x.map_blocks(sparse.COO.from_numpy)

xx = x.mean()
yy = y.mean()

assert_array_equal(xx, yy)
```

```python
Traceback (most recent call last):
  File "dask_sparse.py", line 14, in <module>
    assert_array_equal(xx, yy)
  File "/home/pentschev/.local/lib/python3.5/site-packages/numpy/testing/_private/utils.py", line 894, in assert_array_equal
    verbose=verbose, header='Arrays are not equal')
  File "/home/pentschev/.local/lib/python3.5/site-packages/numpy/testing/_private/utils.py", line 697, in assert_array_compare
    y = array(y, copy=False, subok=True)
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/array/core.py", line 998, in __array__
    x = self.compute()
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/base.py", line 156, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/base.py", line 398, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/threaded.py", line 76, in get
    pack_exception=pack_exception, **kwargs)
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/local.py", line 475, in get_async
    finish(dsk, state, not succeeded)
  File "/usr/lib/python3.5/contextlib.py", line 77, in __exit__
    self.gen.throw(type, value, traceback)
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/callbacks.py", line 99, in local_callbacks
    yield callbacks or ()
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/local.py", line 460, in get_async
    raise_exception(exc, tb)
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/compatibility.py", line 112, in reraise
    raise exc 
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/local.py", line 230, in execute_task
    result = _execute_task(task, data)
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/core.py", line 119, in _execute_task
    return func(*args2)
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/optimization.py", line 942, in __call__
    dict(zip(self.inkeys, args)))
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/core.py", line 149, in get
    result = _execute_task(task, cache)
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/core.py", line 119, in _execute_task
    return func(*args2)
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/compatibility.py", line 93, in apply
    return func(*args, **kwargs)
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/array/reductions.py", line 332, in mean_chunk
    n = numel(x, dtype=dtype, **kwargs)
  File "/home/pentschev/.local/lib/python3.5/site-packages/dask/array/reductions.py", line 323, in numel
    return chunk.sum(np.ones_like(x), **kwargs)
  File "/home/pentschev/.local/lib/python3.5/site-packages/numpy/core/overrides.py", line 151, in public_api
    implementation, public_api, relevant_args, args, kwargs)
  File "/home/pentschev/.local/lib/python3.5/site-packages/numpy/core/numeric.py", line 289, in ones_like
    res = empty_like(a, dtype=dtype, order=order, subok=subok)
  File "/home/pentschev/.local/lib/python3.5/site-packages/numpy/core/overrides.py", line 151, in public_api
    implementation, public_api, relevant_args, args, kwargs)
  File "/home/pentschev/.local/lib/python3.5/site-packages/sparse/sparse_array.py", line 213, in __array__
    raise RuntimeError('Cannot convert a sparse array to dense automatically. '
RuntimeError: Cannot convert a sparse array to dense automatically. To manually densify, use the todense method.
```

Another function that fails with same error besides `mean()` is `std()`. After fixing the error, we have to make sure to add both functions to sparse tests, inside the `functions` list.

Hints from issue discussion:

Commit message: Workaround for mkl.fft failures in test_array_function.py

Patch to review:
diff --git a/.travis.yml b/.travis.yml
index 7f864c96f36..893598788e3 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -47,10 +47,12 @@ jobs:
       - PYTHON=3.7
       - NUMPY=1.16.2
       - PANDAS>=0.24.1
+      - UPSTREAM_SPARSE=1
       - *test_and_lint
       - *no_coverage
       - *no_optimize
       - *imports
+      - *array_function
 
     - env: &py36_dev
       - UPSTREAM_DEV=1  # Install nightly versions of NumPy, pandas, pyarrow
diff --git a/continuous_integration/travis/install.sh b/continuous_integration/travis/install.sh
index 48483888c96..87da7b4ddf3 100644
--- a/continuous_integration/travis/install.sh
+++ b/continuous_integration/travis/install.sh
@@ -75,7 +75,11 @@ if [[ $PYTHONOPTIMIZE != '2' ]] && [[ $NUMPY > '1.11.0' ]] && [[ $NUMPY < '1.14.
 fi
 
 if [[ $NUMPY > '1.13.0' ]]; then
-    pip install sparse 
+    if [[ ${UPSTREAM_SPARSE} ]]; then
+        pip install --upgrade git+https://github.com/pydata/sparse
+    else
+        pip install sparse
+    fi
 fi
 
 if [[ $PYTHON == '2.7' ]]; then
diff --git a/dask/array/reductions.py b/dask/array/reductions.py
index 4d62861517b..7f099e255f4 100644
--- a/dask/array/reductions.py
+++ b/dask/array/reductions.py
@@ -320,7 +320,30 @@ def nanmax(a, axis=None, keepdims=False, split_every=None, out=None):
 
 def numel(x, **kwargs):
     """ A reduction to count the number of elements """
-    return chunk.sum(np.ones_like(x), **kwargs)
+
+    if hasattr(x, 'mask'):
+        return chunk.sum(np.ones_like(x), **kwargs)
+
+    shape = x.shape
+    keepdims = kwargs.get('keepdims', False)
+    axis = kwargs.get('axis', None)
+    dtype = kwargs.get('dtype', np.float64)
+
+    if axis is None:
+        prod = np.prod(shape, dtype=dtype)
+        return np.full((1,) * len(shape), prod, dtype=dtype) if keepdims is True else prod
+
+    if not isinstance(axis, tuple or list):
+        axis = [axis]
+    else:
+        axis = sorted(axis)
+
+    prod = np.prod([shape[dim] for dim in axis])
+    if keepdims is True:
+        new_shape = [shape[dim] if dim not in axis else 1 for dim in range(len(shape))]
+    else:
+        new_shape = [shape[dim] for dim in range(len(shape)) if dim not in axis]
+    return np.full(new_shape, prod, dtype=dtype)
 
 
 def nannumel(x, **kwargs):
diff --git a/dask/array/tests/test_array_function.py b/dask/array/tests/test_array_function.py
index 2ea48c127db..895bf28765b 100644
--- a/dask/array/tests/test_array_function.py
+++ b/dask/array/tests/test_array_function.py
@@ -26,8 +26,6 @@
     lambda x: np.sum(x),
     lambda x: np.var(x),
     lambda x: np.vstack(x),
-    lambda x: np.fft.fft(x.rechunk(x.shape) if isinstance(x, da.Array) else x),
-    lambda x: np.fft.fft2(x.rechunk(x.shape) if isinstance(x, da.Array) else x),
     lambda x: np.linalg.norm(x)])
 def test_array_function_dask(func):
     x = np.random.random((100, 100))
@@ -39,6 +37,19 @@ def test_array_function_dask(func):
     assert_eq(res_y, res_x)
 
 
+@pytest.mark.skipif(missing_arrfunc_cond, reason=missing_arrfunc_reason)
+@pytest.mark.parametrize('func', [np.fft.fft, np.fft.fft2])
+def test_array_function_fft(func):
+    x = np.random.random((100, 100))
+    y = da.from_array(x, chunks=(100, 100))
+    res_x = func(x)
+    res_y = func(y)
+
+    if func.__module__ != 'mkl_fft._numpy_fft':
+        assert isinstance(res_y, da.Array)
+    assert_eq(res_y, res_x)
+
+
 @pytest.mark.skipif(missing_arrfunc_cond, reason=missing_arrfunc_reason)
 @pytest.mark.parametrize('func', [
     lambda x: np.min_scalar_type(x),
diff --git a/dask/array/tests/test_reductions.py b/dask/array/tests/test_reductions.py
index bfc26ec9314..d72bad82cdc 100644
--- a/dask/array/tests/test_reductions.py
+++ b/dask/array/tests/test_reductions.py
@@ -5,6 +5,8 @@
 import pytest
 np = pytest.importorskip('numpy')
 
+import itertools
+
 import dask.array as da
 from dask.array.utils import assert_eq as _assert_eq, same_keys
 from dask.core import get_deps
@@ -15,6 +17,23 @@ def assert_eq(a, b):
     _assert_eq(a, b, equal_nan=True)
 
 
+@pytest.mark.parametrize('dtype', ['f4', 'i4'])
+@pytest.mark.parametrize('keepdims', [True, False])
+def test_numel(dtype, keepdims):
+    x = np.ones((2, 3, 4))
+
+    for length in range(x.ndim):
+        for sub in itertools.combinations([d for d in range(x.ndim)], length):
+            assert_eq(da.reductions.numel(x, axis=sub, keepdims=keepdims, dtype=dtype),
+                      np.sum(x, axis=sub, keepdims=keepdims, dtype=dtype))
+
+    for length in range(x.ndim):
+        for sub in itertools.combinations([d for d in range(x.ndim)], length):
+            ssub = np.random.shuffle(list(sub))
+            assert_eq(da.reductions.numel(x, axis=ssub, keepdims=keepdims, dtype=dtype),
+                      np.sum(x, axis=ssub, keepdims=keepdims, dtype=dtype))
+
+
 def reduction_1d_test(da_func, darr, np_func, narr, use_dtype=True, split_every=True):
     assert_eq(da_func(darr), np_func(narr))
     assert_eq(da_func(darr, keepdims=True), np_func(narr, keepdims=True))
diff --git a/dask/array/tests/test_sparse.py b/dask/array/tests/test_sparse.py
index 459609defab..4d565ffc45f 100644
--- a/dask/array/tests/test_sparse.py
+++ b/dask/array/tests/test_sparse.py
@@ -33,6 +33,12 @@
     lambda x: x.T,
     lambda x: da.transpose(x, (1, 2, 0)),
     lambda x: x.sum(),
+    lambda x: x.mean(),
+    lambda x: x.moment(order=0),
+    pytest.param(lambda x: x.std(),
+                 marks=pytest.mark.xfail(reason="fixed in https://github.com/pydata/sparse/pull/243")),
+    pytest.param(lambda x: x.var(),
+                 marks=pytest.mark.xfail(reason="fixed in https://github.com/pydata/sparse/pull/243")),
     lambda x: x.dot(np.arange(x.shape[-1])),
     lambda x: x.dot(np.eye(x.shape[-1])),
     lambda x: da.tensordot(x, np.ones(x.shape[:2]), axes=[(0, 1), (0, 1)]),


Code review comments:
diff_hunk: @@ -320,7 +320,30 @@ def nanmax(a, axis=None, keepdims=False, split_every=None, out=None):
 
 def numel(x, **kwargs):
     """ A reduction to count the number of elements """
-    return chunk.sum(np.ones_like(x), **kwargs)
+
+    if hasattr(x, 'mask'):
+        return chunk.sum(np.ones_like(x), **kwargs)
+
+    shape = x.shape
+    keepdims = kwargs.get('keepdims', False)
+    axis = kwargs.get('axis', None)
+    dtype = kwargs.get('dtype', np.float64)
+
+    if axis is None:
+        prod = np.prod(shape, dtype=dtype)
+        return np.full((1,) * len(shape), prod, dtype=dtype) if keepdims is True else prod
+
+    if not isinstance(axis, tuple or list):
+        axis = [axis]
+    else:
+        axis = sorted(axis)
+
+    prod = np.prod([shape[dim] for dim in axis])
+    if keepdims is True:
+        new_shape = [shape[dim] if dim not in axis else 1 for dim in range(len(shape))]
+    else:
+        new_shape = [shape[dim] for dim in range(len(shape)) if dim not in axis]
path: dask/array/reductions.py
text: @user1:
These should be `tuple`s since they are shapes.

@author:
Fixed.