how to sample multiple chains in PyMC3

Question

I'm trying to sample multiple chains in PyMC3. In PyMC2 I would do something like this:

for i in range(N):
    model.sample(iter=iter, burn=burn, thin = thin)

How should I do the same thing in PyMC3? I saw there is a 'njobs' argument in the 'sample' method, but it throws an error when I set a value for it. I want to use sampled chains to get 'pymc.gelman_rubin' output.

Answer 1

To run them serially, you can use a similar approach to your PyMC 2 example. The main difference is that each call to sample returns a multi-chain trace instance (containing just a single chain in this case). merge_traces will take a list of multi-chain instances and create a single instance with all the chains.

#!/usr/bin/env python3

import pymc as pm
import numpy as np

from pymc.backends.base import merge_traces

xobs = 4 + np.random.randn(20)

model = pm.Model()
with model:
    mu = pm.Normal('mu', mu=0, sd=20)
    x = pm.Normal('x', mu=mu, sd=1., observed=xobs)
    step = pm.NUTS()

with model:
    trace = merge_traces([pm.sample(1000, step, chain=i)
                          for i in range(2)])

Answer 2

Better is to use njobs to run chains in parallel:

#!/usr/bin/env python3

import pymc3 as pm
import numpy as np

from pymc3.backends.base import merge_traces

xobs = 4 + np.random.randn(20)

model = pm.Model()
with model:
    mu = pm.Normal('mu', mu=0, sd=20)
    x = pm.Normal('x', mu=mu, sd=1., observed=xobs)
    step = pm.NUTS()

with model:
    trace = pm.sample(1000, step, njobs=2)