-
Hello there! Often one would like to sample inside a model. However, sampling requires to split a I had a look at numpyro |
Beta Was this translation helpful? Give feedback.
Replies: 1 comment
-
Have you tried using the lower-level kernel e.g. NUTS and at each step do something like: density_key, sample_key = jax.random.split(rng_key)
logdensity_fn = ft.partial(stochastic_logdensity, density_key)
new_state, info = kernel(sample_key, logdensity_fn, step_size, inverse_mass_matrix) You will need to do that for adaptation as well, and can take inspiration from the window adaptation's implementation. This is slightly more cumbersome than the high-level interface but should work fine. |
Beta Was this translation helpful? Give feedback.
Have you tried using the lower-level kernel e.g. NUTS and at each step do something like:
You will need to do that for adaptation as well, and can take inspiration from the window adaptation's implementation. This is slightly more cumbersome than the high-level interface but should work fine.