Christopher Suter
Christopher Suter
can you share `run_chain` and `construct_hmc`? it's probably a too-large learning rate leading to rejections.
I misspoke. Meant step size not learning rate. Yours is set to 1, which is almost certainly too big. Think of this like the step size parameter in an Euler...
Please use tfp-nightly for JAX support; more discussion here: https://github.com/tensorflow/probability/issues/1994
Can you make sure the old tensorflow-probability pip package isn't still installed?
distrax is installing the non nightly tfp package. This will have to be fixed on the distrax side. Tfp nightly works fine with latest Jax.
Please confirm that you've followed the recommendations in [the release notes](https://github.com/tensorflow/probability/releases/tag/v0.24.0) for keras compatibility, and let us know if that doesn't fix it.