Fixing nested sampling (#1871)

* fixing nested sampling

* first attempt at fixing CI

* import jaxns only if double precision is enabled

* ignore test_nested_sampling in ci

* fix doctest

* implement feedbacks

Author: renecotyfanboy

.github/workflows/ci.yml

@@ -104,7 +104,7 @@ jobs:
     - name: Test with pytest
       run: |
         pytest -vs --durations=20 test/infer/test_mcmc.py
-        pytest -vs --durations=20 test/infer --ignore=test/infer/test_mcmc.py
+        pytest -vs --durations=20 test/infer --ignore=test/infer/test_mcmc.py --ignore=test/contrib/test_nested_sampling.py
         pytest -vs --durations=20 test/contrib --ignore=test/contrib/stochastic_support/test_dcc.py
     - name: Test x64
       run: |
@@ -118,6 +118,9 @@ jobs:
     - name: Test custom prng
       run: |
         JAX_ENABLE_CUSTOM_PRNG=1 pytest -vs test/infer/test_mcmc.py
+    - name: Test nested sampling
+      run: |
+        JAX_ENABLE_X64=1 pytest -vs test/contrib/test_nested_sampling.py
 
 
   examples:

docs/requirements.txt

@@ -4,7 +4,7 @@ funsor
 ipython
 jax
 jaxlib
-jaxns==2.4.8
+jaxns==2.6.3
 Jinja2
 matplotlib
 multipledispatch

docs/source/conf.py

@@ -34,6 +34,8 @@
 
 hmc(None, None)
 
+autodoc_mock_imports = ["jaxns"]
+
 # -- Project information -----------------------------------------------------
 
 project = "NumPyro"

examples/gaussian_shells.py

@@ -121,6 +121,7 @@ def main(args):
 
 if __name__ == "__main__":
     assert numpyro.__version__.startswith("0.15.3")
+
     parser = argparse.ArgumentParser(description="Nested sampler for Gaussian shells")
     parser.add_argument("-n", "--num-samples", nargs="?", default=10000, type=int)
     parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)
@@ -133,6 +134,7 @@ def main(args):
     parser.add_argument("--device", default="cpu", type=str, help='use "cpu" or "gpu".')
     args = parser.parse_args()
 
+    numpyro.enable_x64()
     numpyro.set_platform(args.device)
 
     main(args)

numpyro/contrib/nested_sampling.py

@@ -3,12 +3,13 @@
 
 from functools import singledispatch
 
+import jax
 from jax import random, tree
 import jax.numpy as jnp
 
 try:
+    import jaxns  # noqa: F401
     from jaxns import (
-        DefaultNestedSampler,
         Model,
         Prior,
         TerminationCondition,
@@ -17,11 +18,13 @@
         resample,
         summary,
     )
+    from jaxns.public import DefaultNestedSampler
     from jaxns.utils import NestedSamplerResults
 
 except ImportError as e:
     raise ImportError(
-        "To use this module, please install `jaxns` package. It can be"
+        f"{e} \n "
+        f"To use this module, please install `jaxns>2.5` package. It can be"
         " installed with `pip install jaxns` with python>=3.8"
     ) from e
 
@@ -142,9 +145,7 @@ class NestedSampler:
     :param dict termination_kwargs: keyword arguments to terminate the sampler. Please
         refer to the upstream :meth:`jaxns.NestedSampler.__call__` method.
 
-    **Example**
-
-    .. doctest::
+    Example::
 
         >>> from jax import random
         >>> import jax.numpy as jnp
@@ -258,7 +259,7 @@ def prior_model():
 
         default_constructor_kwargs = dict(
             num_live_points=model.U_ndims * 25,
-            num_parallel_workers=1,
+            devices=jax.devices(),
             max_samples=1e4,
         )
         default_termination_kwargs = dict(dlogZ=1e-4)

setup.py

@@ -61,7 +61,7 @@
             "flax",
             "funsor>=0.4.1",
             "graphviz",
-            "jaxns==2.4.8",
+            "jaxns==2.6.3",
             "matplotlib",
             "optax>=0.0.6",
             "pylab-sdk",  # jaxns dependency

test/contrib/test_nested_sampling.py

@@ -1,6 +1,8 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+import os
+
 import numpy as np
 from numpy.testing import assert_allclose
 import pytest
@@ -11,15 +13,23 @@
 import numpyro
 
 try:
-    from numpyro.contrib.nested_sampling import NestedSampler, UniformReparam
+    if os.environ.get("JAX_ENABLE_X64"):
+        from numpyro.contrib.nested_sampling import NestedSampler, UniformReparam
+
 except ImportError:
     pytestmark = pytest.mark.skip(reason="jaxns is not installed")
+
 import numpyro.distributions as dist
 from numpyro.distributions.transforms import AffineTransform, ExpTransform
 
-pytestmark = pytest.mark.filterwarnings(
-    "ignore:jax.tree_.+ is deprecated:FutureWarning"
-)
+pytestmark = [
+    pytest.mark.filterwarnings("ignore:jax.tree_.+ is deprecated:FutureWarning"),
+    pytest.mark.filterwarnings("ignore:JAX x64"),
+    pytest.mark.skipif(
+        not os.environ.get("JAX_ENABLE_X64"),
+        reason="test suite for jaxns requires double precision",
+    ),
+]
 
 
 # Test helper to extract a few central moments from samples.