Cannot load dumped dok_matrix objects

I had problems trying to load a dok_matrix. Here’s the minimal script to replicate:

from __future__ import print_function
import sys
import joblib
import numpy as np
from scipy.sparse import bsr_matrix, coo_matrix, csc_matrix, csr_matrix, dia_matrix, dok_matrix, lil_matrix

for cls in (bsr_matrix, coo_matrix, csc_matrix, csr_matrix, dia_matrix, lil_matrix, dok_matrix):
    print('Dumping/loading {}'.format(cls))
    m = cls(np.random.normal(size=(100, 100)))
    joblib.dump(m, 'foo')
    joblib.load('foo')

It looks like shape isn’t appropriately set as we load the matrix back out (pickle/cPickle seems to work though). I get the following:

Dumping/loading <class 'scipy.sparse.bsr.bsr_matrix'>
Dumping/loading <class 'scipy.sparse.coo.coo_matrix'>
Dumping/loading <class 'scipy.sparse.csc.csc_matrix'>
Dumping/loading <class 'scipy.sparse.csr.csr_matrix'>
Dumping/loading <class 'scipy.sparse.dia.dia_matrix'>
Dumping/loading <class 'scipy.sparse.lil.lil_matrix'>
Dumping/loading <class 'scipy.sparse.dok.dok_matrix'>
Traceback (most recent call last):
  File "/Users/wradford/Desktop/fail.py", line 11, in <module>
    joblib.load('foo')
  File "/Users/wradford/repos/joblib/joblib/numpy_pickle.py", line 425, in load
    obj = unpickler.load()
  File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/pickle.py", line 858, in load
    dispatch[key](self)
  File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/pickle.py", line 1206, in load_setitems
    dict[stack[i]] = stack[i + 1]
  File "/Users/wradford/repos/joblib/ve/lib/python2.7/site-packages/scipy/sparse/dok.py", line 236, in __setitem__
    0 <= i < self.shape[0] and 0 <= j < self.shape[1]):
  File "/Users/wradford/repos/joblib/ve/lib/python2.7/site-packages/scipy/sparse/base.py", line 499, in __getattr__
    raise AttributeError(attr + " not found")
AttributeError: shape not found

[aabadie]

Digging this up.

Looking at your stacktrace, I think the problem comes from pickle itself. Following http://stackoverflow.com/a/12181159, it seems the problem with matrix_dok comes from the pickle protocol version.

In joblib, the pickle protocol version used is 3 (Default) for python 2 and 4 for python 3 (highest). When calling pickle.dump directly with python 2, the pickle protocol is by default ‘None’ (0) and it’s the only case where the matrix_dok can be reloaded. With any pickle protocol > 0, it fails, hence joblib will fail.

One potential solution is to explicitly set protocol=0 when calling joblib.dump. For instance, the following code works:

from __future__ import print_function
import sys
import joblib
import numpy as np
from scipy.sparse import bsr_matrix, coo_matrix, csc_matrix, csr_matrix, dia_matrix, dok_matrix, lil_matrix

for cls in (bsr_matrix, coo_matrix, csc_matrix, csr_matrix, dia_matrix, lil_matrix, dok_matrix):
    print('Dumping/loading {}'.format(cls))
    m = cls(np.random.normal(size=(100, 100)))
    joblib.dump(m, 'foo', protocol=0)
    joblib.load('foo')

[wejradford]

Thanks guys!