Alloc pickling of GpuArray object

theano/sandbox/gpuarray/tests/GpuArray.pkl

+ctheano.sandbox.gpuarray.type
+GpuArray_unpickler
+p0
+(cnumpy.core.multiarray
+_reconstruct
+p1
+(cnumpy
+ndarray
+p2
+(I0
+tp3
+S'b'
+p4
+tp5
+Rp6
+(I1
+(I1
+tp7
+cnumpy
+dtype
+p8
+(S'f4'
+p9
+I0
+I1
+tp10
+Rp11
+(I3
+S'<'
+p12
+NNNI-1
+I-1
+I0
+tp13
+bI00
+S'\x00\x00(\xc2'
+p14
+tp15
+bNtp16
+Rp17
+.
\ No newline at end of file

theano/sandbox/gpuarray/tests/test_type.py

+import os
+
+from nose.tools import assert_raises
 import numpy
 
 import theano
+from theano.compat import PY3
+from theano import config
 from theano.compile import DeepCopyOp
+from theano.misc.pkl_utils import CompatUnpickler
 
 from .config import test_ctx_name
 from .test_basic_ops import rand_gpuarray
 
 from ..type import GpuArrayType
+from .. import pygpu_activated
+
+if pygpu_activated:
+    import pygpu
 
 
 def test_deep_copy():
@@ -43,3 +53,48 @@ def test_specify_shape():
 def test_filter_float():
     s = theano.shared(numpy.array(0.0, dtype='float32'), target=test_ctx_name)
     theano.function([], updates=[(s, 0.0)])
+
+
+def test_unpickle_gpuarray_as_numpy_ndarray_flag0():
+    oldflag = config.experimental.unpickle_gpu_on_cpu
+    config.experimental.unpickle_gpu_on_cpu = False
+
+    try:
+        testfile_dir = os.path.dirname(os.path.realpath(__file__))
+        fname = 'GpuArray.pkl'
+
+        with open(os.path.join(testfile_dir, fname), 'rb') as fp:
+            if PY3:
+                u = CompatUnpickler(fp, encoding="latin1")
+            else:
+                u = CompatUnpickler(fp)
+            if pygpu_activated:
+                mat = u.load()
+                assert isinstance(mat, pygpu.gpuarray.GpuArray)
+                assert numpy.asarray(mat)[0] == -42.0
+            else:
+                assert_raises(ImportError, u.load)
+    finally:
+        config.experimental.unpickle_gpu_on_cpu = oldflag
+
+
+def test_unpickle_gpuarray_as_numpy_ndarray_flag1():
+    oldflag = config.experimental.unpickle_gpu_on_cpu
+    config.experimental.unpickle_gpu_on_cpu = True
+
+    try:
+        testfile_dir = os.path.dirname(os.path.realpath(__file__))
+        fname = 'GpuArray.pkl'
+
+        with open(os.path.join(testfile_dir, fname), 'rb') as fp:
+            if PY3:
+                u = CompatUnpickler(fp, encoding="latin1")
+            else:
+                u = CompatUnpickler(fp)
+            mat = u.load()
+
+        assert isinstance(mat, numpy.ndarray)
+        assert mat[0] == -42.0
+
+    finally:
+        config.experimental.unpickle_gpu_on_cpu = oldflag

theano/sandbox/gpuarray/type.py

 import numpy
+import six.moves.copyreg as copyreg
+from six import iteritems
+import warnings
 
 import theano
 from theano.tensor.var import _tensor_py_operators
@@ -12,7 +15,7 @@ try:
     from pygpu import gpuarray
     from pygpu.elemwise import compare, elemwise2
 except ImportError:
-    pass
+    pygpu = None
 
 _context_reg = {}
 
@@ -74,7 +77,7 @@ def list_contexts():
 
 # Private method
 def _name_for_ctx(ctx):
-    for k, v in _context_reg:
+    for k, v in iteritems(_context_reg):
         if v == ctx:
             return k
         raise ContextNotDefined('context is not registered')
@@ -751,3 +754,34 @@ Instance of :class:`GpuContextType` to use for the context_type
 declaration of an operation.
 """
 gpu_context_type = GpuContextType()
+
+
+# THIS WORKS But CudaNdarray instances don't compare equal to one
+# another, and what about __hash__ ?  So the unpickled version doesn't
+# equal the pickled version, and the cmodule cache is not happy with
+# the situation.
+def GpuArray_unpickler(npa, ctx_name):
+    if config.experimental.unpickle_gpu_on_cpu:
+        # directly return numpy array
+        warnings.warn(
+            "config.experimental.unpickle_gpu_on_cpu is set to True. "
+            "Unpickling GpuArray as numpy.ndarray")
+        return npa
+    elif pygpu:
+        ctx = get_context(ctx_name)
+        return pygpu.gpuarray.array(npa, copy=True, context=ctx)
+    else:
+        raise ImportError("Cuda not found. Cannot unpickle GpuArray")
+
+copyreg.constructor(GpuArray_unpickler)
+
+
+def GpuArray_pickler(cnda):
+    ctx_name = _name_for_ctx(cnda.context)
+    return (GpuArray_unpickler, (numpy.asarray(cnda), ctx_name))
+
+# In case pygpu is not imported.
+if pygpu is not None:
+    copyreg.pickle(pygpu.gpuarray.GpuArray,
+                   GpuArray_pickler,
+                   GpuArray_unpickler)