Fix unpickling of function with output_keys. (also flake8)

theano/compile/function.py

@@ -10,7 +10,6 @@ import traceback as tb
 import re
 
 from six import string_types
-from theano.compile.io import In
 from theano.compile.function_module import orig_function
 from theano.compile.pfunc import pfunc
 import warnings

theano/compile/function_module.py

@@ -1714,20 +1714,9 @@ class FunctionMaker(object):
         return fn
 
 
-def _pickle_FunctionMaker(self):
-    kwargs = dict(
-        inputs=self.inputs,
-        outputs=self.orig_outputs,
-        fgraph=self.fgraph,
-        mode=self.mode,
-        accept_inplace=self.accept_inplace,
-        function_builder=self.function_builder,
-        profile=self.profile,
-        on_unused_input=self.on_unused_input)
-    return (_constructor_FunctionMaker, (kwargs,))
-
-
 def _constructor_FunctionMaker(kwargs):
+    # Needed for old pickle
+    # Old pickle have at least the problem that output_keys where not safed.
     if theano.config.unpickle_function:
         if theano.config.reoptimize_unpickled_function:
             del kwargs['fgraph']
@@ -1735,8 +1724,6 @@ def _constructor_FunctionMaker(kwargs):
     else:
         return None
 
-copyreg.pickle(FunctionMaker, _pickle_FunctionMaker)
-
 __checkers = []
 
 

theano/compile/tests/test_function_module.py

@@ -649,6 +649,13 @@ class T_picklefunction(unittest.TestCase):
         g(np.asarray(2.))
         self.assertRaises((ValueError, AttributeError), g, 2.)
 
+    def test_output_keys(self):
+        x = T.vector()
+        f = theano.function([x], {'vec': x**2})
+        assert isinstance(f([2, 3, 4]), dict)
+        g = copy.deepcopy(f)
+        assert isinstance(g([2, 3, 4]), dict)
+
     def test_deepcopy_shared_container(self):
         # Ensure that shared containers remain shared after a deep copy.
         a, x = T.scalars('ax')