def _sample_a_lot(rng_key, n, single_sample_fn):
    rng_keys = jax.random.split(rng_key, n)
    return jax.vmap(single_sample_fn)(rng_keys)
def _sample_a_lot(rng_key, n, single_sample_fn):
    rng_keys = jax.random.split(rng_key, n)
    return jax.vmap(single_sample_fn)(rng_keys)