Warmup common or not for 2 chains runs

numpyro
1

Hello,

Bellow a description (Case 1) one may certainly do to perform a 2 chains MCMC sampling. The warm-up of the two chains are independant.
Now, the warmup phase may constitute the large fraction of the time so one may ask of sharing the warmup output state, so I have setup the Case 2.

I was wandering if it is reasonable this Case 2, what would be the prejudice? (nb. For the time beeing the posterior pdf are single mode.)

kernel = NUTS(prob_model,dense_mass=True)

Case 1:

mcmc = MCMC(kernel, num_warmup=5000, num_samples=2000, num_chains=2)
mcmc.run(jax.random.PRNGKey(42), data=data)
samples =  mcmc.get_samples(group_by_chain=True)

Case 2:

mcmc = MCMC(kernel, num_warmup=5000, num_samples=2000, num_chains=1)
keys = random.split(random.PRNGKey(42),4) # keys[0] for later, keys[1] warmup, keys[2].... for independant runs
#
mcmc.warmup(keys[1],data=data)
save_warmup_state =  mcmc.last_state
#
mcmc.post_warmup_state = save_warmup_state
mcmc.run(keys[2],data=data)
samples_1 =  mcmc.get_samples(group_by_chain=True)
#
mcmc.post_warmup_state = save_warmup_state
mcmc.run(keys[3],data=data)
samples_2 =  mcmc.get_samples(group_by_chain=True)
2

Your setup makes sense. It draws 2 chains from a common warmup state. The key point is you are using 2 different keys in run.