Cannot pickle and unpickle vectorized functions.

Problem

I tried to pickle and unpickle an object that contains attributes of vectorized functions that were created with numpy.frompyfunc. The pickle.loads crashed with AttributeError: 'module' object has no attribute 'test (vectorized)'

How to reproduce

I tried to create the minimum example to reproduce the problem and here it is:

import numpy as np

def func(a):
    return a*2

vec_func = np.frompyfunc(func, 1, 1)
pickle.loads(pickle.dumps(vec_func))

This crashes with the following error:

AttributeError                            Traceback (most recent call last)
<ipython-input-33-b5f3d96132c6> in <module>()
      4 
      5 vec_func = np.frompyfunc(func, 1, 1)
----> 6 pickle.loads(pickle.dumps(vec_func))

/usr/lib/python2.7/pickle.pyc in loads(str)
   1386 def loads(str):
   1387     file = StringIO(str)
-> 1388     return Unpickler(file).load()
   1389 
   1390 # Doctest

/usr/lib/python2.7/pickle.pyc in load(self)
    862             while 1:
    863                 key = read(1)
--> 864                 dispatch[key](self)
    865         except _Stop, stopinst:
    866             return stopinst.value

/usr/lib/python2.7/pickle.pyc in load_reduce(self)
   1137         args = stack.pop()
   1138         func = stack[-1]
-> 1139         value = func(*args)
   1140         stack[-1] = value
   1141     dispatch[REDUCE] = load_reduce

/home/sque/workspace/project/venv/local/lib/python2.7/site-packages/numpy/core/__init__.pyc in _ufunc_reconstruct(module, name)
     69     # scipy.special.expit for instance.
     70     mod = __import__(module, fromlist=[name])
---> 71     return getattr(mod, name)
     72 
     73 def _ufunc_reduce(func):

AttributeError: 'module' object has no attribute 'func (vectorized)'

Possible fix

Probably it has to do with the __name__ that frompyfunc sets in the new wrapper function.

[shoyer]

Does it work to pickle functions vectorized with np.vectorize? Those use a different wrapper object…

[arielshulman]

Well, I’ve also encountered this issue. My problem started with np.vectorize which I could not pickle after I’ve used it once. I’ve posted this issue on dill repository, and the great people there helped me to get more information -

1. The problem lays in np.frompyfunc which np.vectorize uses when the otypes=[object].
2. It occurred only when it was used at least once because np.vectorize uses lazy evaluation.
3. I’ve tested my scenario with multiple numpy versions, and it seems it worked on numpy versions prior 1.20.0, but it was an early conclusion since it pickled, but the unpickling didn’t work.
4. It seems the pickling of ufunc is the problem which is created with the np.frompyfunc, and numpy way to pickle it defined here - https://github.com/numpy/numpy/blob/112e63eae5e472a8f8405f68a313d904f7fb5bda/numpy/core/__init__.py#L128-L134 This holds only the function name, which is good to module defined function that are attached to the module using C Python API, like here - https://github.com/numpy/numpy/blob/dd5ab7b11520cf9cfa2129303c02ff30b2960d6f/numpy/core/src/multiarray/methods.c#L2832-L2840 I don’t have much knowledge in Python C API, so I don’t know if there is a way to do this attach of ufunc dynamically to the current module which runs np.frompyfunc. Moreover if you’ll try to get module name of the np.frompyfunc(my_func,1,1) you’ll get None (also with inspect).

P.S. After all said, I’ve found a way to bypass this by creating a VectorizeWrapper which holds the args, kwargs internally and creates the np.vectorize itself like this -

class VectorizeWrapper:
  def __init(self, *args, **kwargs):
     self.args = args
     self.kwargs = kwargs
     self._create_vectorized_function()
  def _create_vectorized_function(self):
    self._vfn = np.vectorize(*self.args, **self.kwargs)
  def __call__(self, v):
    return self._vfn(v)
  def __getstate__(self):
    retrun self.args, self.kwargs
  def __setstate__(self, state):
    self.args, self.kwargs = state
    self._create_vectorized_function()