Keep getting ImportError: cannot import name 'PriorModelGen' from 'jaxns' using numpyro nested sampling example

[mattgrayling]

Describe the bug I have been following the numpyro example for using nested sampling using jaxns. However, I cannot get anywhere with this due to an import error that I keep getting. Specifically, whenever I try this I always get the error ImportError: cannot import name 'PriorModelGen' from 'jaxns'. I've tried isolating this from the rest of the code, just running the import statement by itself but this still doesn't work. I've also tried on multiple different Python versions, on CPU and GPU and got someone else to check who found the same error.

However, I'm aware that this could potentially be an issue with numpyro.contrib.nested_sampling - maybe this will always give an import error because PriorModelGen is not meant to be imported and numpyro instead need to update their code which interfaces with jaxns. If this is the case, please let me know and I will close this issue and raise one in numpyro instead.

Minimal Verifiable Complete Example Running the following short import statement will results in the import error described above.

from jaxns import PriorModelGen

JAXNS version jaxns==2.2.6

[Joshuaalbert]

I think this is because numpyro is slightly behind jaxns's improvements. You can try downgrading jaxns one version at a time until you don't get it.

pip install --ignore-installed jaxns==x.x.x

The releases are here: https://pypi.org/project/jaxns/#history

Let me know which one works.

[mattgrayling]

Thanks! The most recent version that works is 2.1.0, 2.0.1 also works. For other versions there are a variety of unexpected argument errors or import errors which stop it working depending on the version. Do you want me to close this or leave it open until numpyro gets patched?

[Joshuaalbert]

Thanks for reporting. Let's leave this open until numpyro is updated.

[Joshuaalbert]

This is updated in numpyro now