def get_unravel_fn():
state, info, _ = update_loop(init_state, rng_keys[0])
return mcx_ravel_pytree((state, info))
def update_loop(state, key):
keys = jax.random.split(key, num_chains)
state, info = jax.vmap(kernel, in_axes=(0, 0, 0))(keys, parameters, state)
return state, info, mcx_ravel_pytree((state, info))[0]