Repository: numba/numba

Issue:
can't pickle DUFunc objects

```
@vectorize([float64(float64, float64)])
def powerV(x1,x2):
	
	return x1 ** x2

@jit(nopython=True)
def power_mat(m1,m2):

	size1 = m1.size
	size2 = m2.size

	if size1 == size2:		#m1 and m2 are both matrix with same shape or both constant 
		
		retMat = powerV(m1,m2)
		return retMat

	elif size1 == 1 or size2 == 1:		

		retMat = powerV(m1,m2)
		return retMat

	else:
		rowNum1 = m1.shape[0]
		rowNum2 = m2.shape[0]

		if rowNum1 < rowNum2:
			m2 = m2[-rowNum1:]
		elif rowNum1 > rowNum2:
			m1 = m1[-rowNum2:]

		retMat = powerV(m1,m2)
		return retMat
```
the function 'power_mat' cannot pickle DUFunc objects when using the following to parallel
```
Parallel(n_jobs=n_jobs, backend="multiprocessing")(delayed(evalFitTensor)(gp,stockTensor2Eval_,targetMatrix,stockNum,sizeN,fitFunc,windowR,expr_i) for expr_i in gp.evalstr)
```

```
~/anaconda3/lib/python3.6/site-packages/joblib/parallel.py in __call__(self, iterable)
    932 
    933             with self._backend.retrieval_context():
--> 934                 self.retrieve()
    935             # Make sure that we get a last message telling us we are done
    936             elapsed_time = time.time() - self._start_time

~/anaconda3/lib/python3.6/site-packages/joblib/parallel.py in retrieve(self)
    831             try:
    832                 if getattr(self._backend, 'supports_timeout', False):
--> 833                     self._output.extend(job.get(timeout=self.timeout))
    834                 else:
    835                     self._output.extend(job.get())

~/anaconda3/lib/python3.6/multiprocessing/pool.py in get(self, timeout)
    642             return self._value
    643         else:
--> 644             raise self._value
    645 
    646     def _set(self, i, obj):

~/anaconda3/lib/python3.6/multiprocessing/pool.py in _handle_tasks(taskqueue, put, outqueue, pool, cache)
    422                         break
    423                     try:
--> 424                         put(task)
    425                     except Exception as e:
    426                         job, idx = task[:2]

~/anaconda3/lib/python3.6/site-packages/joblib/pool.py in send(obj)
    156             def send(obj):
    157                 buffer = BytesIO()
--> 158                 CustomizablePickler(buffer, self._reducers).dump(obj)
    159                 self._writer.send_bytes(buffer.getvalue())
    160             self._send = send

TypeError: can't pickle DUFunc objects
```

It seems that  the funcion 'powerV' cannot be pickled by joblib for multiprocessing task.  
How can I fix the problem . Thank you. 


Discussion:
if I use the following function to replace, then can be pickled by joblib for multiprocessing task.

```
@jit(nopython=True)
def power_mat(m1,m2):

	size1 = m1.size
	size2 = m2.size

	if size1 == size2:		#m1 and m2 are both matrix with same shape or both constant 
		
		retMat = m1 ** m2
		return retMat

	elif size1 == 1 or size2 == 1:		

		retMat = m1 ** m2
		return retMat

	else:
		rowNum1 = m1.shape[0]
		rowNum2 = m2.shape[0]

		if rowNum1 < rowNum2:
			m2 = m2[-rowNum1:]
		elif rowNum1 > rowNum2:
			m1 = m1[-rowNum2:]

		retMat = m1 ** m2
		return retMat
```

duplication of https://github.com/numba/numba/issues/2943
For #2943, we marked the issue to be a documentation fix with the view that there will be changes to the ufunc machinery but that has not happened.

Looks like the fix is not too difficult to do.  I am going to submit a patch shortly.

Resolving patch:
diff --git a/numba/npyufunc/dufunc.py b/numba/npyufunc/dufunc.py
--- a/numba/npyufunc/dufunc.py
+++ b/numba/npyufunc/dufunc.py
@@ -1,6 +1,6 @@
 from __future__ import absolute_import, print_function, division
 
-import numpy as np
+from numba import serialize
 
 from .. import jit, typeof, utils, types, numpy_support, sigutils
 from ..typing import npydecl
@@ -9,6 +9,7 @@
 from ..dispatcher import Dispatcher
 from .. import array_analysis
 
+
 def make_dufunc_kernel(_dufunc):
     from ..targets import npyimpl
 
@@ -78,19 +79,36 @@ class DUFunc(_internal._DUFunc):
     def __init__(self, py_func, identity=None, cache=False, targetoptions={}):
         if isinstance(py_func, Dispatcher):
             py_func = py_func.py_func
-        self.targetoptions = targetoptions.copy()
-        kws = {}
-        kws['identity'] = ufuncbuilder.parse_identity(identity)
+        dispatcher = jit(target='npyufunc', cache=cache, **targetoptions)(py_func)
+        self._initialize(dispatcher, identity)
 
-        dispatcher = jit(target='npyufunc', cache=cache)(py_func)
-        super(DUFunc, self).__init__(dispatcher, **kws)
+    def _initialize(self, dispatcher, identity):
+        identity = ufuncbuilder.parse_identity(identity)
+        super(DUFunc, self).__init__(dispatcher, identity=identity)
         # Loop over a copy of the keys instead of the keys themselves,
         # since we're changing the dictionary while looping.
         self._install_type()
         self._lower_me = DUFuncLowerer(self)
         self._install_cg()
-        self.__name__ = py_func.__name__
-        self.__doc__ = py_func.__doc__
+        self.__name__ = dispatcher.py_func.__name__
+        self.__doc__ = dispatcher.py_func.__doc__
+
+    def __reduce__(self):
+        siglist = list(self._dispatcher.overloads.keys())
+        return (serialize._rebuild_reduction,
+                (self.__class__, self._dispatcher, self.identity,
+                 self._frozen, siglist))
+
+    @classmethod
+    def _rebuild(cls, dispatcher, identity, frozen, siglist):
+        self = _internal._DUFunc.__new__(cls)
+        self._initialize(dispatcher, identity)
+        # Re-add signatures
+        for sig in siglist:
+            self.add(sig)
+        if frozen:
+            self.disable_compile()
+        return self
 
     def build_ufunc(self):
         """
@@ -98,6 +116,10 @@ def build_ufunc(self):
         """
         return self
 
+    @property
+    def targetoptions(self):
+        return self._dispatcher.targetoptions
+
     @property
     def nin(self):
         return self.ufunc.nin
@@ -283,4 +305,5 @@ def _install_cg(self, targetctx=None):
         targetctx.insert_func_defn(
             [(self._lower_me, self, sig) for sig in (sig0, sig1)])
 
+
 array_analysis.MAP_TYPES.append(DUFunc)
diff --git a/numba/npyufunc/ufuncbuilder.py b/numba/npyufunc/ufuncbuilder.py
--- a/numba/npyufunc/ufuncbuilder.py
+++ b/numba/npyufunc/ufuncbuilder.py
@@ -1,12 +1,10 @@
 # -*- coding: utf-8 -*-
 from __future__ import print_function, division, absolute_import
 
-import warnings
 import inspect
 from contextlib import contextmanager
 
-import numpy as np
-
+from numba import serialize
 from numba.decorators import jit
 from numba.targets.descriptors import TargetDescriptor
 from numba.targets.options import TargetOptions
@@ -41,6 +39,7 @@ def typing_context(self):
     def target_context(self):
         return cpu_target.target_context
 
+
 ufunc_target = UFuncTarget()
 
 
@@ -57,6 +56,22 @@ def __init__(self, py_func, locals={}, targetoptions={}):
         self.locals = locals
         self.cache = NullCache()
 
+    def __reduce__(self):
+        globs = serialize._get_function_globals_for_reduction(self.py_func)
+        return (
+            serialize._rebuild_reduction,
+            (
+                self.__class__,
+                serialize._reduce_function(self.py_func, globs),
+                self.locals, self.targetoptions,
+            ),
+        )
+
+    @classmethod
+    def _rebuild(cls, redfun, locals, targetoptions):
+        return cls(serialize._rebuild_function(*redfun),
+                   locals, targetoptions)
+
     def enable_caching(self):
         self.cache = FunctionCache(self.py_func)
 
@@ -91,7 +106,7 @@ def store_overloads_on_success():
             # use to ensure overloads are stored on success
             try:
                 yield
-            except:
+            except Exception:
                 raise
             else:
                 exists = self.overloads.get(cres.signature)
@@ -131,6 +146,7 @@ def _compile_element_wise_function(nb_func, targetoptions, sig):
     args, return_type = sigutils.normalize_signature(sig)
     return cres, args, return_type
 
+
 def _finalize_ufunc_signature(cres, args, return_type):
     '''Given a compilation result, argument types, and a return type,
     build a valid Numba signature after validating that it doesn't
@@ -146,6 +162,7 @@ def _finalize_ufunc_signature(cres, args, return_type):
     assert return_type != types.pyobject
     return return_type(*args)
 
+
 def _build_element_wise_ufunc_wrapper(cres, signature):
     '''Build a wrapper for the ufunc loop entry point given by the
     compilation result object, using the element-wise signature.
@@ -169,7 +186,8 @@ def _build_element_wise_ufunc_wrapper(cres, signature):
     1: _internal.PyUFunc_One,
     None: _internal.PyUFunc_None,
     "reorderable": _internal.PyUFunc_ReorderableNone,
-    }
+}
+
 
 def parse_identity(identity):
     """
@@ -255,9 +273,11 @@ def build_ufunc(self):
             # For instance, -1 for typenum will cause segfault.
             # If elements of type-list (2nd arg) is tuple instead,
             # there will also memory corruption. (Seems like code rewrite.)
-            ufunc = _internal.fromfunc(self.py_func.__name__, self.py_func.__doc__,
-                                    ptrlist, dtypelist, inct, outct, datlist,
-                                    keepalive, self.identity)
+            ufunc = _internal.fromfunc(
+                self.py_func.__name__, self.py_func.__doc__,
+                ptrlist, dtypelist, inct, outct, datlist,
+                keepalive, self.identity,
+            )
 
             return ufunc
 
@@ -314,9 +334,11 @@ def build_ufunc(self):
         outct = len(self.sout)
 
         # Pass envs to fromfuncsig to bind to the lifetime of the ufunc object
-        ufunc = _internal.fromfunc(self.py_func.__name__, self.py_func.__doc__,
-                                ptrlist, dtypelist, inct, outct, datlist,
-                                keepalive, self.identity, self.signature)
+        ufunc = _internal.fromfunc(
+            self.py_func.__name__, self.py_func.__doc__,
+            ptrlist, dtypelist, inct, outct, datlist,
+            keepalive, self.identity, self.signature,
+        )
         return ufunc
 
     def build(self, cres):


Test patch:
diff --git a/numba/tests/npyufunc/test_dufunc.py b/numba/tests/npyufunc/test_dufunc.py
--- a/numba/tests/npyufunc/test_dufunc.py
+++ b/numba/tests/npyufunc/test_dufunc.py
@@ -1,11 +1,13 @@
 from __future__ import print_function, absolute_import, division
-from numba import unittest_support as unittest
+
+import pickle
 
 import numpy as np
 
-from numba import njit
+from numba import njit, vectorize
 from numba.npyufunc import dufunc
 from ..support import MemoryLeakMixin
+from numba import unittest_support as unittest
 
 
 def pyuadd(a0, a1):
@@ -84,5 +86,40 @@ def test_ufunc_props(self):
         self.assertEqual(duadd.ntypes, 1)
         self.assertEqual(duadd.ntypes, len(duadd.types))
 
+
+class TestDUFuncPickling(MemoryLeakMixin, unittest.TestCase):
+    def check(self, ident, result_type):
+        buf = pickle.dumps(ident)
+        rebuilt = pickle.loads(buf)
+
+        # Check reconstructed dufunc
+        r = rebuilt(123)
+        self.assertEqual(123, r)
+        self.assertIsInstance(r, result_type)
+
+        # Try to use reconstructed dufunc in @jit
+        @njit
+        def foo(x):
+            return rebuilt(x)
+
+        r = foo(321)
+        self.assertEqual(321, r)
+        self.assertIsInstance(r, result_type)
+
+    def test_unrestricted(self):
+        @vectorize
+        def ident(x1):
+            return x1
+
+        self.check(ident, result_type=(int, np.integer))
+
+    def test_restricted(self):
+        @vectorize(["float64(float64)"])
+        def ident(x1):
+            return x1
+
+        self.check(ident, result_type=float)
+
+
 if __name__ == "__main__":
     unittest.main()