Fix issue with pickling Deterministic (#4120)

* Add eight schools pickle test

* Add deterministic test

* adding test for str() explicitly

* Add end of line

* fixes #4112, promote dynamic type to full class

* removing unused import functools

Co-authored-by: Ravin Kumar <[email protected]>

Author: Spaak

pymc3/model.py

@@ -13,7 +13,6 @@
 #   limitations under the License.
 
 import collections
-import functools
 import itertools
 import threading
 import warnings
@@ -1903,14 +1902,22 @@ def _walk_up_rv(rv, formatting='plain'):
     return all_rvs
 
 
-def _repr_deterministic_rv(rv, formatting='plain'):
-    """Make latex string for a Deterministic variable"""
-    if formatting == 'latex':
-        return r"$\text{{{name}}} \sim \text{{Deterministic}}({args})$".format(
-            name=rv.name, args=r",~".join(_walk_up_rv(rv, formatting=formatting)))
-    else:
-        return "{name} ~ Deterministic({args})".format(
-            name=rv.name, args=", ".join(_walk_up_rv(rv, formatting=formatting)))
+class DeterministicWrapper(tt.TensorVariable):
+    def _str_repr(self, formatting='plain'):
+        if formatting == 'latex':
+            return r"$\text{{{name}}} \sim \text{{Deterministic}}({args})$".format(
+                name=self.name, args=r",~".join(_walk_up_rv(self, formatting=formatting)))
+        else:
+            return "{name} ~ Deterministic({args})".format(
+                name=self.name, args=", ".join(_walk_up_rv(self, formatting=formatting)))
+
+    def _repr_latex_(self):
+        return self._str_repr(formatting='latex')
+
+    __latex__ = _repr_latex_
+
+    def __str__(self):
+        return self._str_repr(formatting='plain')
 
 
 def Deterministic(name, var, model=None, dims=None):
@@ -1929,15 +1936,7 @@ def Deterministic(name, var, model=None, dims=None):
     var = var.copy(model.name_for(name))
     model.deterministics.append(var)
     model.add_random_variable(var, dims)
-    var._repr_latex_ = functools.partial(_repr_deterministic_rv, var, formatting='latex')
-    var.__latex__ = var._repr_latex_
-
-    # simply assigning var.__str__ is not enough, since str() will default to the class-
-    # defined __str__ anyway; see https://stackoverflow.com/a/5918210/1692028
-    old_type = type(var)
-    new_type = type(old_type.__name__ + '_pymc3_Deterministic', (old_type,),
-        {'__str__': functools.partial(_repr_deterministic_rv, var, formatting='plain')})
-    var.__class__ = new_type
+    var.__class__ = DeterministicWrapper # adds str and latex functionality
 
     return var
 

pymc3/tests/test_distributions.py

@@ -1774,6 +1774,14 @@ def test___str__(self):
         for str_repr in self.expected_str:
             assert str_repr in model_str
 
+    def test_str(self):
+        for distribution, str_repr in zip(self.distributions, self.expected_str):
+            assert str(distribution) == str_repr
+
+        model_str = str(self.model)
+        for str_repr in self.expected_str:
+            assert str_repr in model_str
+
 
 def test_discrete_trafo():
     with pytest.raises(ValueError) as err:

pymc3/tests/test_model.py

@@ -16,6 +16,7 @@
 import theano
 import theano.tensor as tt
 import numpy as np
+import pickle
 import pandas as pd
 import numpy.testing as npt
 import unittest
@@ -421,3 +422,27 @@ def test_tempered_logp_dlogp():
 
     npt.assert_allclose(func_nograd(x), func(x)[0])
     npt.assert_allclose(func_temp_nograd(x), func_temp(x)[0])
+
+
+def test_model_pickle(tmpdir):
+    """Tests that PyMC3 models are pickleable"""
+    with pm.Model() as model:
+        x = pm.Normal('x')
+        pm.Normal('y', observed=1)
+
+    file_path = tmpdir.join("model.p")
+    with open(file_path, 'wb') as buff:
+        pickle.dump(model, buff)
+
+
+def test_model_pickle_deterministic(tmpdir):
+    """Tests that PyMC3 models are pickleable"""
+    with pm.Model() as model:
+        x = pm.Normal('x')
+        z = pm.Normal("z")
+        pm.Deterministic("w", x/z)
+        pm.Normal('y', observed=1)
+
+    file_path = tmpdir.join("model.p")
+    with open(file_path, 'wb') as buff:
+        pickle.dump(model, buff)