Use dill to serialize logp functions in DensityDist

Author: aseyboldt

RELEASE-NOTES.md

@@ -5,6 +5,7 @@
 ### Maintenance
 - Mentioned the way to do any random walk with `theano.tensor.cumsum()` in `GaussianRandomWalk` docstrings (see [#4048](https://github.com/pymc-devs/pymc3/pull/4048)).
 - Fixed numerical instability in ExGaussian's logp by preventing `logpow` from returning `-inf` (see [#4050](https://github.com/pymc-devs/pymc3/pull/4050)).
+- Use dill to serialize user defined logp functions in `DensityDist`. `dill` is now a required dependency. (see [#3844](https://github.com/pymc-devs/pymc3/issues/3844)).
 
 ### Documentation
 

environment-dev.yml

@@ -37,7 +37,7 @@ dependencies:
   - dataclasses  # python_version < 3.7
   - contextvars  # python_version < 3.7
   - mkl-service
+  - dill
   - libblas=*=*mkl
   - pip:
     - black_nbconvert
-    - dill

pymc3/distributions/distribution.py

@@ -14,6 +14,7 @@
 
 import numbers
 import contextvars
+import dill
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from typing import Optional, Callable
@@ -419,6 +420,19 @@ def __init__(
         self.wrap_random_with_dist_shape = wrap_random_with_dist_shape
         self.check_shape_in_random = check_shape_in_random
 
+    def __getstate__(self):
+        # We use dill to serialize the logp function, as this is almost
+        # always defined in the notebook and won't be pickled correctly.
+        # Fix https://github.com/pymc-devs/pymc3/issues/3844
+        logp = dill.dumps(self.logp)
+        vals = self.__dict__.copy()
+        vals['logp'] = logp
+        return vals
+
+    def __setstate__(self, vals):
+        vals['logp'] = dill.loads(vals['logp'])
+        self.__dict__ = vals
+
     def random(self, point=None, size=None, **kwargs):
         if self.rand is not None:
             not_broadcast_kwargs = dict(point=point)

pymc3/tests/test_distributions.py

@@ -79,6 +79,7 @@
 
 from ..distributions import continuous
 from pymc3.theanof import floatX
+import pymc3 as pm
 from numpy import array, inf, log, exp
 from numpy.testing import assert_almost_equal, assert_allclose, assert_equal
 import numpy.random as nr
@@ -1872,3 +1873,13 @@ def test_issue_3051(self, dims, dist_cls, kwargs):
         assert isinstance(actual_a, np.ndarray)
         assert actual_a.shape == (X.shape[0],)
         pass
+
+
+def test_serialize_density_dist():
+    def func(x):
+        return -2 * (x ** 2).sum()
+
+    with pm.Model():
+        pm.Normal('x')
+        pm.DensityDist('y', func)
+        pm.sample(draws=1, tune=1, mp_ctx="spawn")

requirements-dev.txt

@@ -17,4 +17,4 @@ sphinx-autobuild==0.7.1
 sphinx>=1.5.5
 watermark
 parameterized
-dill
+dill

requirements.txt

@@ -9,3 +9,4 @@ h5py>=2.7.0
 typing-extensions>=3.7.4
 dataclasses; python_version < '3.7'
 contextvars; python_version < '3.7' 
+dill