AttributeError: Can't pickle local object 'Mixture._comp_dist_random_wrapper.<locals>.wrapped_random'

Questions
6

I’m on a Mac, but I’m seeing the same error when trying to pickle a model that contains a mixture where one of the components of the mixture is itself a mixture. Here’s a MWE to reproduce what I’m seeing:

import numpy as np
import pymc3 as pm
import theano.tensor as tt


class MixHelper:

    def __init__(self, pv1, pv2, w):
        self.pv1 = pv1
        self.pv2 = pv2
        self.w = w

    def __call__(self, pt):
        a = self.pv1.logp(pt).sum() + tt.log(self.w[0])
        b = self.pv2.logp(pt).sum() + tt.log(self.w[1])
        return pm.logsumexp([a, b])


def get_model():
    # Number of prior mixture components:
    with pm.Model() as model:
        pv1 = pm.Normal.dist(np.zeros(3), 10., shape=3)

        pv2dists = []
        for k in range(3):
            pvtmp = pm.Normal.dist(np.random.random(size=3), 5., shape=3)
            pv2dists.append(pvtmp)
        pv2 = pm.Mixture.dist(w=np.array([0.6, 0.3, 0.1]),
                            comp_dists=pv2dists,
                            shape=3)

        w = pm.Dirichlet('w', a=np.ones(2))
        mixlogp = MixHelper(pv1, pv2, w)
        v = pm.DensityDist('v', mixlogp, shape=3)

    return model


def main():
    import pickle
    model = get_model()

    with open('test.pkl', 'wb') as f:
        pickle.dump(model, f)

    with open('test.pkl', 'rb') as f:
        model = pickle.load(f)


if __name__ == '__main__':
    main()

Here’s the error I get from running this on macOS 10.15.5, Python v3.7, pymc3 v3.9.2:

Traceback (most recent call last):
File "simple_test.py", line 61, in <module>
    main()
File "simple_test.py", line 47, in main
    pickle.dump(model, f)
AttributeError: Can't pickle local object 'Mixture._comp_dist_random_wrapper.<locals>.wrapped_random'