Repository: dask/dask

Issue:
Wrapper always pass None to random.choice
From dask/array/random.py:

>     with ignoring(AttributeError):
>         @doc_wraps(np.random.RandomState.choice)
>         def choice(self, a, size=None, replace=True, p=None, chunks=None):
>             return self._wrap(np.random.RandomState.choice, a,
>                               size=size, replace=True, p=None, chunks=chunks)

Wrapper gets non-None p argument -> Wrapper passes None to p kwarg of np.random.RandomState.choice anyway.

Test:

> import dask
> import dask.array as da
> import numpy as np
> import dask.array.random as dar
> import matplotlib.pyplot as plt
> 
> p = np.arange(12)
> p = p*1.0 / np.sum(p)
> 
> p = da.from_array(p, chunks=6)
> s = dar.RandomState()
> 
> c = s.choice(len(p), p=p,size=10000,chunks=8)
> cc = c.compute()
> b = np.bincount(cc)
> print(b)
> 
> c = np.random.choice(len(p), p=np.array(p), size=10000)
> b = np.bincount(c)
> print(b)
> 

Result from numpy and dask are different.

> [806 815 833 873 877 809 850 818 837 792 892 798]
> [   0  117  332  467  637  731  901 1079 1251 1316 1483 1686]



Discussion:

Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index 4a28b01fe..50c066574 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -2497,8 +2497,8 @@ def broadcast_shapes(*shapes):
         return shapes[0]
     out = []
     for sizes in zip_longest(*map(reversed, shapes), fillvalue=-1):
-        dim = max(sizes)
-        if any(i != -1 and i != 1 and i != dim and not np.isnan(i) for i in sizes):
+        dim = 0 if 0 in sizes else max(sizes)
+        if any(i not in [-1, 0, 1, dim] and not np.isnan(i) for i in sizes):
             raise ValueError("operands could not be broadcast together with "
                              "shapes {0}".format(' '.join(map(str, shapes))))
         out.append(dim)
diff --git a/dask/array/random.py b/dask/array/random.py
index ad5ef1954..7ce9d2549 100644
--- a/dask/array/random.py
+++ b/dask/array/random.py
@@ -1,11 +1,12 @@
 from __future__ import absolute_import, division, print_function
 
 from itertools import product
+from numbers import Integral
 from operator import getitem
 
 import numpy as np
 
-from .core import (normalize_chunks, Array, slices_from_chunks,
+from .core import (normalize_chunks, Array, slices_from_chunks, asarray,
                    broadcast_shapes, broadcast_to)
 from .. import sharedict
 from ..base import tokenize
@@ -170,8 +171,62 @@ class RandomState(object):
     with ignoring(AttributeError):
         @doc_wraps(np.random.RandomState.choice)
         def choice(self, a, size=None, replace=True, p=None, chunks=None):
-            return self._wrap(np.random.RandomState.choice, a,
-                              size=size, replace=True, p=None, chunks=chunks)
+            dsks = []
+            # Normalize and validate `a`
+            if isinstance(a, Integral):
+                # On windows the output dtype differs if p is provided or
+                # absent, see https://github.com/numpy/numpy/issues/9867
+                dummy_p = np.array([1]) if p is not None else p
+                dtype = np.random.choice(1, size=(), p=dummy_p).dtype
+                len_a = a
+                if a < 0:
+                    raise ValueError("a must be greater than 0")
+            else:
+                a = asarray(a).rechunk(a.shape)
+                dtype = a.dtype
+                if a.ndim != 1:
+                    raise ValueError("a must be one dimensional")
+                len_a = len(a)
+                dsks.append(a.dask)
+                a = a._keys()[0]
+
+            # Normalize and validate `p`
+            if p is not None:
+                if not isinstance(p, Array):
+                    # If p is not a dask array, first check the sum is close
+                    # to 1 before converting.
+                    p = np.asarray(p)
+                    if not np.isclose(p.sum(), 1, rtol=1e-7, atol=0):
+                        raise ValueError("probabilities do not sum to 1")
+                    p = asarray(p)
+                else:
+                    p = p.rechunk(p.shape)
+
+                if p.ndim != 1:
+                    raise ValueError("p must be one dimensional")
+                if len(p) != len_a:
+                    raise ValueError("a and p must have the same size")
+
+                dsks.append(p.dask)
+                p = p._keys()[0]
+
+            if size is None:
+                size = ()
+            elif not isinstance(size, (tuple, list)):
+                size = (size,)
+
+            chunks = normalize_chunks(chunks, size)
+            sizes = list(product(*chunks))
+            state_data = random_state_data(len(sizes), self._numpy_state)
+
+            name = 'da.random.choice-%s' % tokenize(state_data, size, chunks,
+                                                    a, replace, p)
+            keys = product([name], *(range(len(bd)) for bd in chunks))
+            dsk = {k: (_choice, state, a, size, replace, p) for
+                   k, state, size in zip(keys, state_data, sizes)}
+
+            return Array(sharedict.merge((name, dsk), *dsks),
+                         name, chunks, dtype=dtype)
 
     # @doc_wraps(np.random.RandomState.dirichlet)
     # def dirichlet(self, alpha, size=None, chunks=None):
@@ -352,6 +407,11 @@ class RandomState(object):
                           size=size, chunks=chunks)
 
 
+def _choice(state_data, a, size, replace, p):
+    state = np.random.RandomState(state_data)
+    return state.choice(a, size=size, replace=replace, p=p)
+
+
 def _apply_random(func, state_data, size, args, kwargs):
     """Apply RandomState method with seed"""
     state = np.random.RandomState(state_data)
@@ -368,6 +428,8 @@ seed = _state.seed
 beta = _state.beta
 binomial = _state.binomial
 chisquare = _state.chisquare
+if hasattr(_state, 'choice'):
+    choice = _state.choice
 exponential = _state.exponential
 f = _state.f
 gamma = _state.gamma
diff --git a/dask/array/routines.py b/dask/array/routines.py
index ed1ed2482..edb76bfcb 100644
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -14,7 +14,7 @@ from toolz import concat, sliding_window, interleave
 
 from .. import sharedict
 from ..core import flatten
-from ..base import tokenize, compute_as_if_collection
+from ..base import tokenize
 from . import numpy_compat, chunk
 
 from .core import (Array, map_blocks, elemwise, from_array, asarray,
@@ -530,11 +530,18 @@ def round(a, decimals=0):
 
 @wraps(np.unique)
 def unique(x):
-    name = 'unique-' + x.name
-    dsk = {(name, i): (np.unique, key) for i, key in enumerate(x.__dask_keys__())}
-    parts = compute_as_if_collection(Array, sharedict.merge((name, dsk), x.dask),
-                                     list(dsk.keys()))
-    return np.unique(np.concatenate(parts))
+    x = x.ravel()
+
+    out = atop(np.unique, "i", x, "i", dtype=x.dtype)
+    out._chunks = tuple((np.nan,) * len(c) for c in out.chunks)
+
+    name = 'unique-aggregate-' + out.name
+    dsk = {(name, 0): (np.unique, (np.concatenate, out._keys()))}
+    out = Array(
+        sharedict.merge((name, dsk), out.dask), name, ((np.nan,),), out.dtype
+    )
+
+    return out
 
 
 @wraps(np.roll)
diff --git a/dask/dataframe/io/io.py b/dask/dataframe/io/io.py
index e2d0a9633..7c5942eba 100644
--- a/dask/dataframe/io/io.py
+++ b/dask/dataframe/io/io.py
@@ -240,7 +240,7 @@ def from_bcolz(x, chunksize=None, categorize=True, index=None, lock=lock,
                     np.issubdtype(x.dtype[name], np.unicode_) or
                     np.issubdtype(x.dtype[name], np.object_)):
                 a = da.from_array(x[name], chunks=(chunksize * len(x.names),))
-                categories[name] = da.unique(a)
+                categories[name] = da.unique(a).compute()
 
     columns = tuple(x.dtype.names)
     divisions = tuple(range(0, len(x), chunksize))
diff --git a/dask/diagnostics/profile_visualize.py b/dask/diagnostics/profile_visualize.py
index dee158e7b..b66691432 100644
--- a/dask/diagnostics/profile_visualize.py
+++ b/dask/diagnostics/profile_visualize.py
@@ -237,6 +237,7 @@ def plot_tasks(results, dsk, palette='Viridis', label_size=60, **kwargs):
 
     defaults = dict(title="Profile Results",
                     tools="hover,save,reset,xwheel_zoom,xpan",
+                    toolbar_location='above',
                     plot_width=800, plot_height=300)
     defaults.update((k, v) for (k, v) in kwargs.items() if k in
                     _get_figure_keywords())
@@ -317,6 +318,7 @@ def plot_resources(results, palette='Viridis', **kwargs):
 
     defaults = dict(title="Profile Results",
                     tools="save,reset,xwheel_zoom,xpan",
+                    toolbar_location='above',
                     plot_width=800, plot_height=300)
     defaults.update((k, v) for (k, v) in kwargs.items() if k in
                     _get_figure_keywords())
@@ -324,15 +326,18 @@ def plot_resources(results, palette='Viridis', **kwargs):
         t, mem, cpu = zip(*results)
         left, right = min(t), max(t)
         t = [i - left for i in t]
-        p = bp.figure(y_range=(0, max(cpu)), x_range=(0, right - left), **defaults)
+        p = bp.figure(y_range=fix_bounds(0, max(cpu), 100),
+                      x_range=fix_bounds(0, right - left, 1),
+                      **defaults)
     else:
         t = mem = cpu = []
-        p = bp.figure(y_range=(0, 100), x_range=(0, 10), **defaults)
+        p = bp.figure(y_range=(0, 100), x_range=(0, 1), **defaults)
     colors = palettes.all_palettes[palette][6]
     p.line(t, cpu, color=colors[0], line_width=4, legend='% CPU')
     p.yaxis.axis_label = "% CPU"
-    p.extra_y_ranges = {'memory': Range1d(start=(min(mem) if mem else 0),
-                                          end=(max(mem) if mem else 100))}
+    p.extra_y_ranges = {'memory': Range1d(*fix_bounds(min(mem) if mem else 0,
+                                                      max(mem) if mem else 100,
+                                                      100))}
     p.line(t, mem, color=colors[2], y_range_name='memory', line_width=4,
            legend='Memory')
     p.add_layout(LinearAxis(y_range_name='memory', axis_label='Memory (MB)'),
@@ -341,6 +346,11 @@ def plot_resources(results, palette='Viridis', **kwargs):
     return p
 
 
+def fix_bounds(start, end, min_span):
+    """Adjust end point to ensure span of at least `min_span`"""
+    return start, max(end, start + min_span)
+
+
 def plot_cache(results, dsk, start_time, metric_name, palette='Viridis',
                label_size=60, **kwargs):
     """Visualize the results of profiling in a bokeh plot.
@@ -374,6 +384,7 @@ def plot_cache(results, dsk, start_time, metric_name, palette='Viridis',
 
     defaults = dict(title="Profile Results",
                     tools="hover,save,reset,wheel_zoom,xpan",
+                    toolbar_location='above',
                     plot_width=800, plot_height=300)
     defaults.update((k, v) for (k, v) in kwargs.items() if k in
                     _get_figure_keywords())
diff --git a/docs/source/array-api.rst b/docs/source/array-api.rst
index d3c760b01..b6e18ff5e 100644
--- a/docs/source/array-api.rst
+++ b/docs/source/array-api.rst
@@ -228,6 +228,7 @@ Random
    random.beta
    random.binomial
    random.chisquare
+   random.choice
    random.exponential
    random.f
    random.gamma
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index 62253d869..f10ec3098 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -11,6 +11,8 @@ Array
 - Add ``allclose`` (:pr:`2771`)
 - Remove ``random.different_seeds`` from Dask Array API docs (:pr:`2772`)
 - Deprecate ``vnorm`` in favor of ``dask.array.linalg.norm`` (:pr:`2773`)
+- Reimplement ``unique`` to be lazy (:pr:`2775`)
+- Support broadcasting of Dask Arrays with 0-length dimensions (:pr:`2784`)
 
 DataFrame
 +++++++++
@@ -37,6 +39,7 @@ Core
 Array
 +++++
 
+-  ``da.random.choice`` now works with array arguments (:pr:`2781`)
 -  Support indexing in arrays with np.int (fixes regression) (:pr:`2719`)
 -  Handle zero dimension with rechunking (:pr:`2747`)
 -  Support -1 as an alias for "size of the dimension" in ``chunks`` (:pr:`2749`)


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 141a65cba..29bbce220 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -407,6 +407,7 @@ def test_binops():
 
 
 def test_broadcast_shapes():
+    assert (0, 5) == broadcast_shapes((0, 1), (1, 5))
     assert (3, 4, 5) == broadcast_shapes((3, 4, 5), (4, 1), ())
     assert (3, 4) == broadcast_shapes((3, 1), (1, 4), (4,))
     assert (5, 6, 7, 3, 4) == broadcast_shapes((3, 1), (), (5, 6, 7, 1, 4))
@@ -574,7 +575,7 @@ def test_broadcast_to():
     x = np.random.randint(10, size=(5, 1, 6))
     a = from_array(x, chunks=(3, 1, 3))
 
-    for shape in [a.shape, (5, 4, 6), (2, 5, 1, 6), (3, 4, 5, 4, 6)]:
+    for shape in [a.shape, (5, 0, 6), (5, 4, 6), (2, 5, 1, 6), (3, 4, 5, 4, 6)]:
         xb = chunk.broadcast_to(x, shape)
         ab = broadcast_to(a, shape)
 
@@ -590,7 +591,7 @@ def test_broadcast_to():
 def test_broadcast_to_array():
     x = np.random.randint(10, size=(5, 1, 6))
 
-    for shape in [(5, 4, 6), (2, 5, 1, 6), (3, 4, 5, 4, 6)]:
+    for shape in [(5, 0, 6), (5, 4, 6), (2, 5, 1, 6), (3, 4, 5, 4, 6)]:
         a = np.broadcast_to(x, shape)
         d = broadcast_to(x, shape)
 
@@ -600,13 +601,35 @@ def test_broadcast_to_array():
 def test_broadcast_to_scalar():
     x = 5
 
-    for shape in [tuple(), (2, 3), (5, 4, 6), (2, 5, 1, 6), (3, 4, 5, 4, 6)]:
+    for shape in [tuple(), (0,), (2, 3), (5, 4, 6), (2, 5, 1, 6), (3, 4, 5, 4, 6)]:
         a = np.broadcast_to(x, shape)
         d = broadcast_to(x, shape)
 
         assert_eq(a, d)
 
 
+@pytest.mark.parametrize('u_shape, v_shape', [
+    [tuple(), (2, 3)],
+    [(1,), (2, 3)],
+    [(1, 1), (2, 3)],
+    [(0, 3), (1, 3)],
+    [(2, 0), (2, 1)],
+    [(1, 0), (2, 1)],
+    [(0, 1), (1, 3)],
+])
+def test_broadcast_operator(u_shape, v_shape):
+    u = np.random.random(u_shape)
+    v = np.random.random(v_shape)
+
+    d_u = from_array(u, chunks=1)
+    d_v = from_array(v, chunks=1)
+
+    w = u * v
+    d_w = d_u * d_v
+
+    assert_eq(w, d_w)
+
+
 @pytest.mark.parametrize('original_shape,new_shape,chunks', [
     ((10,), (10,), (3, 3, 4)),
     ((10,), (10, 1, 1), 5),
diff --git a/dask/array/tests/test_random.py b/dask/array/tests/test_random.py
index 38e0a73c6..fa25de9ef 100644
--- a/dask/array/tests/test_random.py
+++ b/dask/array/tests/test_random.py
@@ -216,3 +216,52 @@ def test_multinomial():
         y = np.random.multinomial(20, [1 / 6.] * 6, size=size)
 
         assert x.shape == y.shape == x.compute().shape
+
+
+def test_choice():
+    np_dtype = np.random.choice(1, size=()).dtype
+    size = (10, 3)
+    chunks = 4
+    x = da.random.choice(3, size=size, chunks=chunks)
+    assert x.dtype == np_dtype
+    assert x.shape == size
+    res = x.compute()
+    assert res.dtype == np_dtype
+    assert res.shape == size
+
+    np_a = np.array([1, 3, 5, 7, 9], dtype='f8')
+    da_a = da.from_array(np_a, chunks=2)
+
+    for a in [np_a, da_a]:
+        x = da.random.choice(a, size=size, chunks=chunks)
+        res = x.compute()
+        assert x.dtype == np_a.dtype
+        assert res.dtype == np_a.dtype
+        assert set(np.unique(res)).issubset(np_a)
+
+    np_p = np.array([0, 0.2, 0.2, 0.3, 0.3])
+    da_p = da.from_array(np_p, chunks=2)
+
+    for a, p in [(da_a, np_p), (np_a, da_p)]:
+        x = da.random.choice(a, size=size, chunks=chunks, p=p)
+        res = x.compute()
+        assert x.dtype == np_a.dtype
+        assert res.dtype == np_a.dtype
+        assert set(np.unique(res)).issubset(np_a[1:])
+
+    np_dtype = np.random.choice(1, size=(), p=np.array([1])).dtype
+    x = da.random.choice(5, size=size, chunks=chunks, p=np_p)
+    res = x.compute()
+    assert x.dtype == np_dtype
+    assert res.dtype == np_dtype
+
+    errs = [(-1, None),             # negative a
+            (np_a[:, None], None),  # a must be 1D
+            (np_a, np_p[:, None]),  # p must be 1D
+            (np_a, np_p[:-2]),      # a and p must match
+            (3, np_p),              # a and p must match
+            (4, [0.2, 0.2, 0.3])]   # p must sum to 1
+
+    for (a, p) in errs:
+        with pytest.raises(ValueError):
+            da.random.choice(a, size=size, chunks=chunks, p=p)
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
index 68ce4c332..f91d1cfe6 100644
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -473,9 +473,15 @@ def test_round():
 
 
 def test_unique():
-    x = np.array([1, 2, 4, 4, 5, 2])
-    d = da.from_array(x, chunks=(3,))
-    assert_eq(da.unique(d), np.unique(x))
+    a = np.array([1, 2, 4, 4, 5, 2])
+    d = da.from_array(a, chunks=(3,))
+
+    r_a = np.unique(a)
+    r_d = da.unique(d)
+
+    assert isinstance(r_d, da.Array)
+
+    assert_eq(r_d, r_a)
 
 
 def _maybe_len(l):
diff --git a/dask/diagnostics/tests/test_profiler.py b/dask/diagnostics/tests/test_profiler.py
index b42a85fdc..5c062640b 100644
--- a/dask/diagnostics/tests/test_profiler.py
+++ b/dask/diagnostics/tests/test_profiler.py
@@ -245,13 +245,21 @@ def test_resource_profiler_plot():
     assert len(p.tools) == 1
     assert isinstance(p.tools[0], bokeh.models.HoverTool)
     assert check_title(p, "Not the default")
-    # Test empty, checking for errors
-    rprof.clear()
-
-    with pytest.warns(None) as record:
-        rprof.visualize(show=False, save=False)
 
-    assert len(record) == 0
+    # Test with empty and one point, checking for errors
+    rprof.clear()
+    for results in [[], [(1.0, 0, 0)]]:
+        rprof.results = results
+        with pytest.warns(None) as record:
+            p = rprof.visualize(show=False, save=False)
+        assert len(record) == 0
+        # Check bounds are valid
+        assert p.x_range.start == 0
+        assert p.x_range.end == 1
+        assert p.y_range.start == 0
+        assert p.y_range.end == 100
+        assert p.extra_y_ranges['memory'].start == 0
+        assert p.extra_y_ranges['memory'].end == 100
 
 
 @pytest.mark.skipif("not bokeh")