bump jaxns to >=2.0.1

[Joshuaalbert]

• reqs are now jaxns>==2.0.1
• adjusted the wrapped to use new structure.
• Note, that jaxns likes the user to not jit-compile the top-level nested sampling run. This is because it breaks up the algorithm into static parts and non-static parts.

EDIT: I changed JAXNS to support 3.8, so the bumped version is to >=2.0.1 now.

[Joshuaalbert]

@fehiepsi here you go^. I haven't tested it. I'm not sure if you have a test for the wrapper. Otherwise, would you mind running the Gaussian Shells example in numpyro.

[fehiepsi]

Woohoo, thanks @Joshuaalbert! We have some tests here - I can take a look later today if things work.

[fehiepsi]

Thanks @Joshuaalbert! The changes look great to me. After fixing a typo in PriorModel definition (missing Prior(...)), the tests passed and gaussian shell example also passed in my system.

You will need to run make format (with flake8, black, isort installed) to fix formating issues.

[Joshuaalbert]

I notice that the test on 3.8 if failing. That's because jaxns requires 3.9+.

ERROR: Ignored the following versions that require a different python version: 2.0.0 Requires-Python >=3.9
ERROR: Could not find a version that satisfies the requirement jaxns>=2.0.0 (from versions: 0.0.1, 0.0.2, 0.0.3, 0.0.4, 0.0.5, 0.0.6, 0.0.7, 1.0.0, 1.0.1, 1.1.0, 1.1.1, 1.1.2)

[fehiepsi]

LGTM, just have some nits left 💯

[Joshuaalbert]

Nits done

[fehiepsi]

jaxns requires 3.9+

I guess we can just specify jaxns and skip the test on CI by replacing this line by

try:
    from numpyro.contrib.nested_sampling import NestedSampler, UniformReparam
except ImportError:
    pytestmark = pytest.mark.skip(reason="jaxns is not installed")

For the docs, please feel free to bump the version to 3.9 at this line.

[Joshuaalbert]

Looks like there will still be a problem though because the setup.py will try to get jaxns>=2.0.0 this will fail before getting to the tests. Edit: The only solution would be to put a little logic in setup.py so that it only tries to install jaxns when the python version >= 3.9. This plus skipping the pytest would work.

[Joshuaalbert]

@fehiepsi as expected it still fails. The only option would be for me to add a conditional check to include jaxns if python >= 3.9.

[Joshuaalbert]

@fehiepsi fixed the E402. You can run again.

[fehiepsi]

the current issue is unrelated. let me fit it upstream

[fehiepsi]

Thanks, @Joshuaalbert! It turns out that there are missing libraries in jaxns requirements like matplotlib and pylab-sdk. Could you merge the dev branch? Hopefully it will work now.