I cannot figure out how to implement a NUTS based simulated annealing algorithm to approximate the maximum of an arbitrary potential function. More specifically, the potential function requires a sequential computation therefore it needs to be jitted. In order to run a simulated annealing optimizer, there needs to exist a manually controlled temperature parameter to the potential function, whereas the rest of the parameters needs to be sampled via MCMC. Moreover, this temperature parameter can be adjusted at each iteration, and the acceptance ratio should be calculated using the latest temperature value. How can I achieve so?

`infer.MCMC`

samplers will be able to use it in the way you want. Instead, could you look into jax’s more direct stochastic optimizers? Maybe use the `SVI`

module, which is a sort of middle ground between the two?