Choice of sampler + additional args to log-density call

[homerjed]

Hi all,

I'm here as one of the types of users "People who just need state-of-the-art samplers that are fast, robust and well tested".

I'm looking to sample a very high-dimensional posterior (image sized) with the use of some sharding (multi-device) and a log-density function that uses a probabilistic estimate of the log-likelihood. This means I need to pass a jax.random.PRNGKey to my log-density call in the sampler.

Could I get some pointers on the best samplers for this task, and possibly how to throw in additional args to my log-density call in the kernel?

Thank you!

[homerjed]

the MCLMC example works fast on my setup with a 20,000-dimensional problem.