Custom Sampler with MH or Rejection steps

is there any way of having a sampler like NUTS that at random iterations will make proposals and accept them based on some function?
Basically, I would need two functions: the proposal and the acceptance. I have seen two closely related samplers, one is a discussion of the impossibility of a general Metropolis Hasting algorithm, the other one is the HMCGibs sampler.
To make it more concrete, one particular sampler could use NUTS and after warmup, at every iteration, it choses two random variables and flip their value if the energy is increased.
Is there any way of implementing what I want with the current samplers?

Thanks !

Should something like that be implemented with “postprocess_fn” ? What are the inputs of postprocess_fn ? I see (self, model_args, model_kwargs), can I get the potential function into my wrapper for the evaluation ?

Thanks !

The inputs would be model arguments. I think you can access potential function using this attribute. There are two issues:

  • it would be tricky to provide random keys in that method
  • this only modifies the collected results that you want to obtain. If you want to do modify the proposal for the next MCMC step, you need to wrap the sample method.

Thanks for the answer !
Ok, so the postprocess_fn is not the way to go.
What sample method are you referring to? If I am using HMC, NUTS, they only look at the log_likelihood of any distribution, they do not look at the sample method, am I right?
Best !

It is the kernel’s sample method. You can check this document: Markov Chain Monte Carlo (MCMC) — NumPyro documentation