I assume you put the seed handler outside of vmap? The vmap body function should be closed, i.e. meaning that it does not depend on some global seed handler. You can do something like vmap(lambda key, sample: log_density(seed(guide, key), …))(keys, samples) where keys = random.split(random.PRNGKey(0), num_samples)