Repository: dask/dask

Issue:
ResourceProfile has different behaviour using context manager vs `register`
I'm using the dask profilers with the pattern
```
for profiler in profilers: profiler.register()
foo() # Do work, some of it dask, some of it other
for profiler in profilers: profiler.unregister()
```
(although with some try/finally stuff to make it more robust). I had assumed that this would be equivalent to running `foo` under a context manager per profiler, except that the `with` statement doesn't allow for unpacking so can't take a variable number of context managers (and `contextlib.nested` is deprecated and I'm on Python 2 so can't use `ExitStack`). However, the `ResourceProfiler` code has specific checks for whether it is currently being used as a context manager (`_enter` member) and if not, it stops tracking resources as soon as a computation is complete (in `_finish`), which leads to gaps in the graph whenever no dask computation is running.

Example code:
```python
#!/usr/bin/env python
import dask.array as da
from dask.diagnostics import ResourceProfiler
import time

def foo():
    x = da.ones(1000000000, chunks=1000000)
    da.dot(x, x).compute()  # Keep things busy for a while
    time.sleep(3)
    da.dot(x, x).compute()  # Keep things busy again

profiler = ResourceProfiler(dt=0.2)
profiler.register()
foo()
profiler.unregister()
profiler.visualize()
```

Discussion:
jcrist: @bmerry does this clarification make sense?

Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index 3f3f222db..3446e5723 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -2451,20 +2451,25 @@ def take(a, indices, axis=0):
 
 @wraps(np.compress)
 def compress(condition, a, axis=None):
+    from .wrap import zeros
+
     if axis is None:
-        raise NotImplementedError("Must select axis for compression")
+        a = a.ravel()
+        axis = 0
     if not -a.ndim <= axis < a.ndim:
         raise ValueError('axis=(%s) out of bounds' % axis)
     if axis < 0:
         axis += a.ndim
 
-    condition = np.array(condition, dtype=bool)
+    condition = asarray(condition).astype(bool)
     if condition.ndim != 1:
         raise ValueError("Condition must be one dimensional")
     if len(condition) < a.shape[axis]:
-        condition = condition.copy()
-        condition.resize(a.shape[axis])
+        condition_xtr = zeros(a.shape[axis], dtype=bool, chunks=a.chunks[axis])
+        condition_xtr = condition_xtr[len(condition):]
+        condition = concatenate([condition, condition_xtr])
 
+    condition = np.array(condition)
     slc = ((slice(None),) * axis + (condition, ) +
            (slice(None),) * (a.ndim - axis - 1))
     return a[slc]
diff --git a/dask/diagnostics/profile.py b/dask/diagnostics/profile.py
index ad41e189e..41543c868 100644
--- a/dask/diagnostics/profile.py
+++ b/dask/diagnostics/profile.py
@@ -134,6 +134,10 @@ class ResourceProfiler(Callback):
     manually.
 
     >>> prof.clear()  # doctest: +SKIP
+
+    Note that when used as a context manager data will be collected throughout
+    the duration of the enclosed block. In contrast, when registered globally
+    data will only be collected while a dask scheduler is active.
     """
     def __init__(self, dt=1):
         self._tracker = _Tracker(dt)


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 61934ac1b..04d9ba159 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -432,15 +432,17 @@ def test_compress():
     x = np.arange(25).reshape((5, 5))
     a = from_array(x, chunks=(2, 2))
 
+    assert_eq(np.compress([True, False, True, False, True], x),
+              da.compress([True, False, True, False, True], a))
     assert_eq(np.compress([True, False, True, False, True], x, axis=0),
               da.compress([True, False, True, False, True], a, axis=0))
     assert_eq(np.compress([True, False, True, False, True], x, axis=1),
               da.compress([True, False, True, False, True], a, axis=1))
+    assert_eq(np.compress([True, False], x),
+              da.compress([True, False], a))
     assert_eq(np.compress([True, False], x, axis=1),
               da.compress([True, False], a, axis=1))
 
-    with pytest.raises(NotImplementedError):
-        da.compress([True, False], a)
     with pytest.raises(ValueError):
         da.compress([True, False], a, axis=100)
     with pytest.raises(ValueError):