Skip to content

Refactor (Jax) Metropolis sampler to avoid using jax.experimental.loops

Created by: PhilipVinc

According to https://github.com/google/jax/pull/10278 jax.experimental.loops is scheduled for deprecation in a few months. I had written netket's metropolis sampler using it.

We should update the body of the sampler to avoid using the experimental functionality. It should be relatively straightforward.

This is the statement that must be updated https://github.com/netket/netket/blob/a15bf8796fb91bb430d0d6079c16a3d0f69e5472/netket/sampler/metropolis.py#L347