Hi @smp, if you only want to be jit-compiled (i.e. not include grad, vmap, pmap transformations), then you might want to use host_callback. Some of its examples in numpyro:
- to display progress_bar in mcmc
- to implement samplers for truncated distributions