Hi @fonnesbeck, thanks for the reply!
I tried to use sample_numpyro_nuts, installing Jax and numpyro. However, I get the following exception: Exception: Disabling of omnistaging is no longer supported in JAX version 0.2.12 and higher. When I use an earlier version of Jax, then I get AttributeError: module 'jaxlib.xla_client' has no attribute 'get_local_backend'. Apparently other people have the same problem (here).
May I ask you how you get numpyro nuts to work? Thank you and sorry if this is slightly distant from the original question.