Skip to content

jax samplers: allow passing rng_key via kwargs

Vicentini Filippo requested to merge github/fork/inailuig/jax-metropolislocal-rngkey into master

Created by: inailuig

Currently there is no way to pass a rng_key to the jax version of nk.sampler.MetropolisLocal. This PR tries to add that functionality.

The same could also be done for the other samplers (MetropolisExchange, ...)

Merge request reports