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)