How to use TFPKernel in NumPyro?

I would like to use TFPKernel with NumPyro, so according to the documentation I wrote it like this:

kernel = numpyro.infer.TFPKernel[tfp.mcmc.ReplicaExchangeMC](model)

But got the following error.

AttributeError: module ‘numpyro.infer’ has no attribute ‘TFPKernel’

I would appreciate it if someone could tell me how to access the TFP kernel.

You can access it through this path: numpyro.contrib.tfp.mcmc — NumPyro documentation Probably you are the first user of this module. Let us know if you have any feedback. We intended to deprecate this to reduce maintaining, but if there are usage cases, we’ll try to keep it as-is.

Thank you! I was able to access the class successfully.
The function of ReplicaExchangeMC is important to our research, and this wrap is one of the main reasons we chose NumPyro.
I will happily give you feedbacks later.

The wrapped TFP replica exchange MCMC had following problems.
(1) Sample parameters do not move at all.
(2) There is no way to know each output data belongs to which inverse temperature.

If possible, I would like to implement replica exhange MCMC on original NumPyro.
The easiest way is pause runnning MCMC, exchange parameters according to inverse temperature between chains, and resume.

Could you please answer the following questions if possible?

  • If rerunning the mcmc object, is the posterior distribution inherited or discarded? Do I need to joint them each rerun?
  • Is there a way to know every posterior and sample parameters of each chain?
  • Is it possible to set parameters such as inverse temperature for each chain?
1 Like