Fixing nested sampling

.github/workflows/ci.yml

@@ -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.2
 Jinja2
 matplotlib
 multipledispatch

examples/gaussian_shells.py

@@ -18,6 +18,8 @@
        https://github.com/joshspeagle/dynesty/blob/master/demos/Examples%20--%20Gaussian%20Shells.ipynb
 """
 
+# ruff: noqa: E402
+
 import argparse
 
 import matplotlib.pyplot as plt
@@ -30,6 +32,9 @@
 import numpyro.distributions as dist
 from numpyro.infer import MCMC, NUTS, DiscreteHMCGibbs
 
+numpyro.enable_x64()
+numpyro.set_host_device_count(2)
+
 
 class GaussianShell(dist.Distribution):
     support = dist.constraints.real_vector

numpyro/contrib/nested_sampling.py

@@ -3,12 +3,13 @@
 
 from functools import singledispatch
 
+import jax
 from jax import random
 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
 
@@ -258,7 +261,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)

numpyro/distributions/transforms.py

@@ -719,7 +719,7 @@ class LowerCholeskyAffine(Transform):
        >>> scale_tril = jnp.array([[0.3, 0.0], [1.0, 0.5]])
        >>> affine = LowerCholeskyAffine(loc=loc, scale_tril=scale_tril)
        >>> affine(base)
-       Array([0.3, 1.5], dtype=float32)
+       Array([0.3, 1.5], dtype=float64)
     """
 
     domain = constraints.real_vector

setup.py

@@ -61,7 +61,7 @@
             "flax",
             "funsor>=0.4.1",
             "graphviz",
-            "jaxns==2.4.8",
+            "jaxns==2.6.2",
             "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
@@ -17,9 +19,13 @@
 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.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.