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

Questions
7

Ah, I see now — this happens because of the way Mixture._comp_dist_random_wrapper is implemented, which prevents defining a model within a function or method. So one option for me is to move the model definition out of the function. An alternate fix for my use case was to get rid of pm.Mixture.dist() and use my custom mixture class instead:

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


class MixHelper:

    def __init__(self, K, w, ps):
        self.K = K
        self.ps = ps
        self.w = w

    def __call__(self, pt):
        logps = []
        for k in range(self.K):
            logp = self.ps[k].logp(pt).sum() + tt.log(self.w[k])
            logps.append(logp)
        return pm.logsumexp(logps)


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.DensityDist.dist(
            MixHelper(K=3, w=np.array([0.6, 0.3, 0.1]), ps=pv2dists),
            shape=3)

        w = pm.Dirichlet('w', a=np.ones(2))
        v = pm.DensityDist('v', MixHelper(K=2, w=w, ps=[pv1, pv2]), 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()