Fixing nested sampling

[renecotyfanboy]

There are some import discrepancies between jaxns<2.6 and jaxns 2.6.*. This small fix seems to do the trick!

[fehiepsi]

could you update jaxns version in setup and docs/requirements

numpyro/setup.py

Line 64 in 8e9313f

"jaxns==2.4.8",

[renecotyfanboy]

@fehiepsi any idea on how to enforce double precision in the test base for this specific test ? I don't know how JAX and pytest work together and if it is possible to enable x64 for a single test

[fehiepsi]

You can add a flag like this one

numpyro/.github/workflows/ci.yml

Line 111 in 8e9313f

JAX_ENABLE_X64=1 pytest -vs test/infer/test_mcmc.py -k x64

[renecotyfanboy]

Okay I hope it works now, I had a lot of trouble since latest versions of jaxns enforces double precision to summarize, I had to :

• Makes a special case for pytest if double precision is not enabled to skip NS related tests
• Add an extra invoke for dedicated jaxns tests
• Change the expected result of a doctest in affine transform in the documentation, as the output is jnp.float64 instead of float32. All the other tests ran smoothly on my computer
• Updated a bit the jaxns contributed code. To be fair, at some point I should take some time to update it properly with the new interfaces proposed by the package

[renecotyfanboy]

Okay @fehiepsi this should work now. For you to know, most of the struggles came from the fact that importing jaxns automatically enable double precision. So top-level imports in pytests triggered the double precision even if not required. I don't import it if the double precision is disabled, and the related test suite is automatically skipped in this situation.

[renecotyfanboy]

I think it's ready for review

[fehiepsi]

Woohoo, thanks @renecotyfanboy!