Repository: dask/dask

Issue:
Dask take fails with error in __bool__
I have an issue with dask's implementation of `take`. I expect it to work as numpy's `take`, but it fails to compute. See the following snippet for an example:

``` python
def test_dask_take():
    x = np.arange(10.0, 0, -1)
    choices = np.asarray([0, 0, 1, 1, 2, 2, 1, 1, 0, 0])
    res = x.take(choices)
    np.testing.assert_allclose([10., 10., 9., 9., 8., 8., 9., 9., 10., 10.], res)

    x_da = da.from_array(x, chunks=(x.shape[0]))
    choices_da = da.from_array(choices, chunks=(x.shape[0]))
    res2 = da.take(x_da, choices_da, axis=0).compute()  # <-- fails

    np.testing.assert_allclose(res, res2)
```

The traceback from py.test for this failure:

```
    ================================== FAILURES ===================================
    _______________________________ test_dask_take ________________________________

        def test_dask_take():
            x = np.arange(10.0, 0, -1)
            choices = np.asarray([0, 0, 1, 1, 2, 2, 1, 1, 0, 0])
            res = x.take(choices)
            np.testing.assert_allclose([10., 10., 9., 9., 8., 8., 9., 9., 10., 10.]
    , res)

            x_da = da.from_array(x, chunks=(x.shape[0]))
            choices_da = da.from_array(choices, chunks=(x.shape[0]))
    >       res2 = da.take(x_da, choices_da, axis=0).compute()

    test_dask.py:202:
    _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
    ../dask/dask/array/core.py:1301: in take
        return a[(slice(None),) * axis + (indices,)]
    ../dask/dask/array/core.py:750: in __getitem__
        dsk, chunks = slice_array(out, self.name, self.chunks, index)
    ../dask/dask/array/slicing.py:100: in slice_array
        if all(index == slice(None, None, None) for index in index):
    _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

    self = dask.array<x_4, shape=(10,), chunks=((10,)), dtype=bool>

        def __bool__(self):
    >       return bool(self.compute())
    E       ValueError: The truth value of an array with more than one element is a
    mbiguous. Use a.any() or a.all()

    ../dask/dask/array/core.py:722: ValueError
```



Discussion:
Dask.array.take currently expects a list or numpy array for the choices argument, not another dask.array.   We can index an out-of-core array with an in-core array, but not an out-of-core array with another out-of-core array.

It is hard to do this with another dask array because we are unable to look at the values of the choices dask.array to determine which input blocks affect which output blocks.

So this would work 

``` python
>>> da.take(x_da, choices, axis=0).compute()
array([ 10.,  10.,   9.,   9.,   8.,   8.,   9.,   9.,  10.,  10.])
```

But this is actually a fairly hard problem to do generally:

``` python
>>> da.take(x_da, choices_da, axis=0).compute()
```

This page on [dask.array's limitations](http://dask.readthedocs.org/en/latest/array-overview.html#limitations) might be a good read.

It could be that what you want to do is achievable through other means.  Why were you using `da.take`?

Agreed.

`da.take` with two dask arrays would be complex. However, I think that `da.take` could support a second case. Think of selecting lots of `choices` from a small `numpy.ndarray` `x`:

``` python
>>> da.take(x, choices_da, axis=0).compute()
```

A small numpy array `x` can be held available, when evaluating choices on every chunk of `choices_da`. I.e.  `da.take(DA, NP)` is about taking a few samples from a big array, whereas `da.take(NP, DA)` is about choosing lots of values from a small sample.

Yes, I agree that that case is probably possible.  

Is this something you are actively trying to use or something that you've noted is possible (I'm trying to arrange this in my list of priorities.)

This is something I am actively trying to use, not just an observation along the way.


Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -1402,13 +1402,24 @@ def concatenate(seq, axis=0):
 
 
 @wraps(np.take)
-def take(a, indices, axis):
+def take(a, indices, axis=0):
     if not -a.ndim <= axis < a.ndim:
         raise ValueError('axis=(%s) out of bounds' % axis)
     if axis < 0:
         axis += a.ndim
-    return a[(slice(None),) * axis + (indices,)]
+    if isinstance(a, np.ndarray) and isinstance(indices, Array):
+        return _take_dask_array_from_numpy(a, indices, axis)
+    else:
+        return a[(slice(None),) * axis + (indices,)]
+
+
+def _take_dask_array_from_numpy(a, indices, axis):
+    assert isinstance(a, np.ndarray)
+    assert isinstance(indices, Array)
 
+    return indices.map_blocks(lambda block: np.take(a, block, axis),
+                              chunks=indices.chunks,
+                              dtype=a.dtype)
 
 @wraps(np.transpose)
 def transpose(a, axes=None):


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -1027,7 +1027,6 @@ def test_map_blocks3():
 def test_numpy_compat_is_notimplemented():
     a = np.arange(10)
     x = da.from_array(a, chunks=5)
-
     assert raises(NotImplementedError, lambda: x + a)
 
 
@@ -1043,3 +1042,13 @@ def test_cache():
     assert len(z.dask) < 6  # very short graph
     assert z.chunks == y.chunks
     assert eq(y, z)
+
+
+def test_take_dask_from_numpy():
+    x = np.arange(5).astype('f8')
+    y = da.from_array(np.array([1, 2, 3, 3, 2 ,1]), chunks=3)
+
+    z = da.take(x * 2, y)
+
+    assert z.chunks == y.chunks
+    assert eq(z, np.array([2., 4., 6., 6., 4., 2.]))