Repository: dask/dask

Issue:
Cannot pickle dask array objects
**What happened**:

Attempting to pickle a dask array gives:
```
AttributeError: Can't pickle local object 'broadcast_trick.<locals>.inner'
```

Dask arrays are no longer pickleable since v2.23.0.

**What you expected to happen**:

The dask array can be pickled.

**Minimal Complete Verifiable Example**:

```python
import dask.array
import pickle
a = dask.array.ones((10, 10))
pickle.dumps(a)
```

This fails on versions larger than v2.22.

**Anything else we need to know?**:

I am happy to fix this. It seems like a simple fix to move the `inner` function to the module level.

**Environment**:

- Dask version: 2.25.0
- Python version: 3.7.7
- Operating System: Mac OS X
- Install method (conda, pip, source): pip


Discussion:

Patch:
diff --git a/dask/array/wrap.py b/dask/array/wrap.py
index 4c7be06a0..9a319953b 100644
--- a/dask/array/wrap.py
+++ b/dask/array/wrap.py
@@ -135,6 +135,11 @@ def wrap(wrap_func, func, **kwargs):
 w = wrap(wrap_func_shape_as_first_arg)
 
 
+@curry
+def _broadcast_trick_inner(func, shape, *args, **kwargs):
+    return np.broadcast_to(func((), *args, **kwargs), shape)
+
+
 def broadcast_trick(func):
     """
     Provide a decorator to wrap common numpy function with a broadcast trick.
@@ -155,8 +160,7 @@ def broadcast_trick(func):
     so should be safe.
     """
 
-    def inner(shape, *args, **kwargs):
-        return np.broadcast_to(func((), *args, **kwargs), shape)
+    inner = _broadcast_trick_inner(func)
 
     if func.__doc__ is not None:
         inner.__doc__ = func.__doc__


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 197159398..9d7700e7b 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -3298,12 +3298,18 @@ def test_from_array_names():
     assert set(names.values()) == set([1, 5])
 
 
-def test_array_picklable():
+@pytest.mark.parametrize(
+    "array",
+    [
+        da.arange(100, chunks=25),
+        da.ones((10, 10), chunks=25),
+    ],
+)
+def test_array_picklable(array):
     from pickle import loads, dumps
 
-    a = da.arange(100, chunks=25)
-    a2 = loads(dumps(a))
-    assert_eq(a, a2)
+    a2 = loads(dumps(array))
+    assert_eq(array, a2)
 
 
 def test_from_array_raises_on_bad_chunks():