Repository: dask/dask

Issue:
da.where breaks on non-bool condition
Using dask 0.15.1, numpy 1.13.1, Python 2.7, the following code throws an exception with dask, while the equivalent numpy code works correctly:
```python
>>> import numpy as np
>>> import dask.array as da
>>> a = da.arange(4, dtype=np.uint8, chunks=(4,))
>>> da.where(a, 1, 0).compute()
Traceback (most recent call last):
  File "dask-where.py", line 6, in <module>
    da.where(a, 1, 0).compute()
  File "/home/bmerry/work/sdp/env/local/lib/python2.7/site-packages/dask/base.py", line 98, in compute
    (result,) = compute(self, traverse=False, **kwargs)
  File "/home/bmerry/work/sdp/env/local/lib/python2.7/site-packages/dask/base.py", line 205, in compute
    results = get(dsk, keys, **kwargs)
  File "/home/bmerry/work/sdp/env/local/lib/python2.7/site-packages/dask/threaded.py", line 75, in get
    pack_exception=pack_exception, **kwargs)
  File "/home/bmerry/work/sdp/env/local/lib/python2.7/site-packages/dask/local.py", line 521, in get_async
    raise_exception(exc, tb)
  File "/home/bmerry/work/sdp/env/local/lib/python2.7/site-packages/dask/local.py", line 290, in execute_task
    result = _execute_task(task, data)
  File "/home/bmerry/work/sdp/env/local/lib/python2.7/site-packages/dask/local.py", line 271, in _execute_task
    return func(*args2)
  File "/home/bmerry/work/sdp/env/local/lib/python2.7/site-packages/dask/array/core.py", line 2796, in variadic_choose
    return np.choose(a, choices)
  File "/home/bmerry/work/sdp/env/local/lib/python2.7/site-packages/numpy/core/fromnumeric.py", line 354, in choose
    return _wrapfunc(a, 'choose', choices, out=out, mode=mode)
  File "/home/bmerry/work/sdp/env/local/lib/python2.7/site-packages/numpy/core/fromnumeric.py", line 57, in _wrapfunc
    return getattr(obj, method)(*args, **kwds)
ValueError: invalid entry in choice array
```

From what I can see, `da.where` is being implemented using `np.choose` rather than `np.where`, and so `np.choose` is trying to use the raw value of the condition rather than its truthiness to index one of the two outputs.

Discussion:
jakirkham: Looks much cleaner than the original code. Thanks for making this change.

Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index c880fa0cd..3f3f222db 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -2842,12 +2842,10 @@ def where(condition, x=None, y=None):
     x = broadcast_to(x, shape).astype(dtype)
     y = broadcast_to(y, shape).astype(dtype)
 
-    if isinstance(condition, (bool, np.bool8)):
-        if condition:
-            return x
-        else:
-            return y
+    if np.isscalar(condition):
+        return x if condition else y
     else:
+        condition = asarray(condition).astype('bool')
         return choose(condition, [y, x])
 
 
@@ -3955,6 +3953,11 @@ def repeat(a, repeats, axis=None):
     if not isinstance(repeats, Integral):
         raise NotImplementedError("Only integer valued repeats supported")
 
+    if -a.ndim <= axis < 0:
+        axis += a.ndim
+    elif not 0 <= axis <= a.ndim - 1:
+        raise ValueError("axis(=%d) out of bounds" % axis)
+
     if repeats == 1:
         return a
 


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 701d2674a..61934ac1b 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -686,11 +686,15 @@ def test_choose():
 
 def test_where():
     x = np.random.randint(10, size=(15, 16))
+    x[5, 5] = x[4, 4] = 0 # Ensure some false elements
     d = from_array(x, chunks=(4, 5))
     y = np.random.randint(10, size=15).astype(np.uint8)
     e = from_array(y, chunks=(4,))
 
     for c1, c2 in [(d > 5, x > 5),
+                   (d, x),
+                   (1, 1),
+                   (0, 0),
                    (True, True),
                    (np.True_, np.True_),
                    (False, False),
@@ -698,7 +702,6 @@ def test_where():
         for b1, b2 in [(0, 0), (-e[:, None], -y[:, None])]:
             w1 = where(c1, d, b1)
             w2 = np.where(c2, x, b2)
-
             assert_eq(w1, w2)
 
 
@@ -708,7 +711,7 @@ def test_where_bool_optimization():
     y = np.random.randint(10, size=(15, 16))
     e = from_array(y, chunks=(4, 5))
 
-    for c in [True, False, np.True_, np.False_]:
+    for c in [True, False, np.True_, np.False_, 1, 0]:
         w1 = where(c, d, e)
         w2 = np.where(c, x, y)
 
@@ -2764,7 +2767,7 @@ def test_repeat():
     d = da.from_array(x, chunks=(4, 5, 3))
 
     repeats = [1, 2, 5]
-    axes = [0, 1, 2]
+    axes = [-3, -2, -1, 0, 1, 2]
 
     for r in repeats:
         for a in axes:
@@ -2781,6 +2784,10 @@ def test_repeat():
     with pytest.raises(NotImplementedError):
         da.repeat(d, 2)
 
+    for invalid_axis in [3, -4]:
+        with pytest.raises(ValueError):
+            da.repeat(d, 2, axis=invalid_axis)
+
     x = np.arange(5)
     d = da.arange(5, chunks=(2,))